Update app.py

app.py

@@ -1,48 +1,76 @@
 import gradio as gr
+
 import torch
-from transformers import pipeline, PegasusForConditionalGeneration, PegasusTokenizer
+from transformers import pipeline
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
 
-# Q&A model
-qa_classifier = pipeline(
-    "question-answering",
+classifier = pipeline(
+    "question-answering", 
     model="deepset/roberta-base-squad2",
     tokenizer="deepset/roberta-base-squad2"
 )
 
-# paraphrase model
-paraphrase_model_name = 'tuner007/pegasus_paraphrase'
-paraphrase_tokenizer = PegasusTokenizer.from_pretrained(paraphrase_model_name)
-paraphrase_model = PegasusForConditionalGeneration.from_pretrained(paraphrase_model_name)
+model_name = 'tuner007/pegasus_paraphrase'
 torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
+tokenizer3 = PegasusTokenizer.from_pretrained(model_name)
+model3 = PegasusForConditionalGeneration.from_pretrained(model_name).to(torch_device)
+
+
+def paraphrase(text_input):
+    question = "What is the Sri Lanka's rank on human development Index?"
+    prediction = classifier(
+        context=text_input,
+        question=question,
+        truncation=True,
+        max_length=512,
+        padding=True,
+    )
+    answer = prediction['answer']
+    answer_start = prediction['start']
+    answer_end = prediction['end']
+    context = text_input.split(".")
+    for i in range(len(context)):
+        if answer in context[i]:
+            sentence = context[i].strip() + "."
+            break
+    batch = tokenizer3([sentence],truncation=True,padding='longest',max_length=60, return_tensors="pt").to(torch_device)
+    translated = model3.generate(**batch,max_length=60,num_beams=10, num_return_sequences=1, temperature=1.5)
+    paraphrase = tokenizer3.batch_decode(translated, skip_special_tokens=True)[0]
+    return paraphrase
+
 
-def find_answer_and_paraphrase(text_input, question):
-    # Find the answer
-    answer = qa_classifier(context=text_input, question=question)["answer"]
-    
-    # Find the sentence containing the answer
-    sentences = text_input.split(".")
-    sentence_with_answer = ""
-    for sentence in sentences:
-        if answer in sentence:
-            sentence_with_answer = sentence.strip() + "."
+def qa_paraphrase(text_input):
+    question = "What is the Sri Lanka's rank on human development Index?"
+    prediction = classifier(
+        context=text_input,
+        question=question,
+        truncation=True,
+        max_length=512,
+        padding=True,
+    )
+    answer = prediction['answer']
+    answer_start = prediction['start']
+    answer_end = prediction['end']
+    context = text_input.split(".")
+    for i in range(len(context)):
+        if answer in context[i]:
+            sentence = context[i].strip() + "."
             break
-    
-    # Paraphrase the sentence containing the answer
-    input_text = "paraphrase: " + sentence_with_answer
-    num_return_sequences = 1
-    num_beams = 5
-    batch = paraphrase_tokenizer([input_text], truncation=True, padding='longest', max_length=60, return_tensors="pt").to(torch_device)
-    translated = paraphrase_model.generate(**batch, max_length=60, num_beams=num_beams, num_return_sequences=num_return_sequences, temperature=1.5)
-    paraphrase_text = paraphrase_tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
-
-    return {"Answer": answer, "Paraphrase": paraphrase_text}
-
-
-inputs = [
-    gr.inputs.Textbox(label="Text Input"),
-    gr.inputs.Textbox(label="Question")
-]
-outputs = [gr.outputs.Textbox(label="Answer"), gr.outputs.Textbox(label="Paraphrase")]
-
-iface = gr.Interface(fn=find_answer_and_paraphrase, inputs=inputs, outputs=outputs, title="Question Answering and Paraphrasing", description="Enter the text and question to find the answer and paraphrase the sentence containing the answer.")
-iface.launch()
+    batch = tokenizer3([sentence],truncation=True,padding='longest',max_length=60, return_tensors="pt").to(torch_device)
+    translated = model3.generate(**batch,max_length=60,num_beams=10, num_return_sequences=1, temperature=1.5)
+    paraphrase = tokenizer3.batch_decode(translated, skip_special_tokens=True)[0]
+    return {"Answer": answer, "Paraphrased Sentence": paraphrase}
+
+
+iface = gr.Interface(
+    fn=qa_paraphrase,
+    inputs=gr.inputs.Textbox(label="Text Input"),
+    outputs=[
+        gr.outputs.Textbox(label="Answer"),
+        gr.outputs.Textbox(label="Paraphrased Sentence")
+    ],
+    title="Question Answering and Paraphrasing",
+    description="Answer a given question from a given text and paraphrase the sentence containing the answer."
+)
+
+iface.launch()