Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
+from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline, AutoModelForSeq2SeqLM
 from pdfminer.high_level import extract_text
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
@@ -10,10 +10,15 @@ from rank_bm25 import BM25Okapi
 
 nltk.download('punkt')
 
-model_name = "deepset/roberta-large-squad2"
-model = AutoModelForQuestionAnswering.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer)
+# QA model
+qa_model_name = "deepset/roberta-large-squad2"
+qa_model = AutoModelForQuestionAnswering.from_pretrained(qa_model_name)
+qa_tokenizer = AutoTokenizer.from_pretrained(qa_model_name)
+qa_pipeline = pipeline("question-answering", model=qa_model, tokenizer=qa_tokenizer)
+
+# Summarization model
+summarization_model_name = "facebook/bart-large-cnn"
+summarizer = pipeline("summarization", model=summarization_model_name)
 
 def read_pdf(file):
     try:
@@ -48,34 +53,50 @@ def answer_question(pdf, question, num_words):
         response = qa_pipeline(question=question, context=relevant_text)
         answer = response['answer']
 
-        words = answer.split()
-        if len(words) >= num_words:
-            answer = " ".join(words[:num_words])  # Strict truncation
+        answer = answer.strip()
+        answer = " ".join(answer.split())
+
+        if len(answer.split()) > num_words:
+            try:
+                summarized_answer = summarizer(answer, max_length=num_words+10, min_length=1) # increased max length for better summaries
+                answer = summarized_answer[0]['summary_text']
+                answer = answer.strip()
+                answer = " ".join(answer.split())
+                if len(answer.split()) > num_words:
+                    answer = " ".join(answer.split()[:num_words]) #truncate if summary is still too long
+            except RuntimeError as e:
+                if "Input length of input_ids is" in str(e) and "but `max_length` is set to" in str(e): # more robust check for context window error
+                    answer = " ".join(answer.split()[:num_words])
+                else:
+                    return f"Summarization Error: {e}"
+            except Exception as e:
+                return f"Summarization Error: {e}"
         elif relevant_text: #only add context if there is relevant text
-            remaining_words = num_words - len(words)
+            remaining_words = num_words - len(answer.split())
             added_words = 0
             for sentence in sent_tokenize(relevant_text):
                 sentence_words = sentence.split()
                 words_to_add = min(remaining_words - added_words, len(sentence_words))
-                words.extend(sentence_words[:words_to_add])
+                answer_words = answer.split()
+                answer_words.extend(sentence_words[:words_to_add])
+                answer = " ".join(answer_words)
                 added_words += words_to_add
                 if added_words == remaining_words:
                     break
-            answer = " ".join(words)
-            if len(answer.split()) > num_words: #truncate again if needed
+            if len(answer.split()) > num_words:
                 answer = " ".join(answer.split()[:num_words])
         return answer.strip()
 
     except Exception as e:
         return str(e)
 
-with gr.Blocks() as iface: #use blocks for custom layout
-    gr.Markdown("PDF Q&A with RoBERTa | Made by NP")
+with gr.Blocks() as iface:
+    gr.Markdown("PDF Q&A with RoBERTa")
     with gr.Row():
         pdf_input = gr.File(type="filepath", label="Upload PDF")
         question_input = gr.Textbox(lines=2, placeholder="Ask a question")
         num_words_slider = gr.Slider(minimum=1, maximum=500, value=100, step=1, label="Number of Words")
-    answer_output = gr.Textbox(label="Answer", lines=5) # increased lines for better display
+    answer_output = gr.Textbox(label="Answer", lines=5)
     btn = gr.Button("Submit")
     btn.click(fn=answer_question, inputs=[pdf_input, question_input, num_words_slider], outputs=answer_output)