ethanrom/chat2

1import gradio as gr2 3import torch4from transformers import pipeline5from transformers import PegasusForConditionalGeneration, PegasusTokenizer6 7classifier = pipeline(8    "question-answering", 9    model="deepset/roberta-base-squad2",10    tokenizer="deepset/roberta-base-squad2"11)12 13model_name = 'tuner007/pegasus_paraphrase'14torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'15tokenizer3 = PegasusTokenizer.from_pretrained(model_name)16model3 = PegasusForConditionalGeneration.from_pretrained(model_name).to(torch_device)17 18 19def qa_paraphrase(text_input, question):20    prediction = classifier(21        context=text_input,22        question=question,23        truncation=True,24        max_length=512,25        padding=True,26    )27    answer = prediction['answer']28    answer_start = prediction['start']29    answer_end = prediction['end']30    context = text_input.split(".")31    for i in range(len(context)):32        if answer in context[i]:33            sentence = context[i].strip() + "."34            break35    batch = tokenizer3([sentence],truncation=True,padding='longest',max_length=60, return_tensors="pt").to(torch_device)36    translated = model3.generate(**batch,max_length=60,num_beams=10, num_return_sequences=1, temperature=1.5)37    paraphrase = tokenizer3.batch_decode(translated, skip_special_tokens=True)[0]38    return f"Answer: {answer}\nLong Form Answer: {paraphrase}"39 40 41iface = gr.Interface(42    fn=qa_paraphrase,43    inputs=[44        gr.inputs.Textbox(label="Text Input"),45        gr.inputs.Textbox(label="Question")46    ],47    outputs=gr.outputs.Textbox(label="Output"),48    title="Long Form Question Answering",49    description="mimics long form question answering by extracting the sentence containing the answer and paraphrasing it"50)51 52iface.launch()53