Update app.py

app.py

@@ -6,12 +6,11 @@ from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 import nltk
 from nltk.tokenize import sent_tokenize
-from rank_bm25 import BM25Okapi  # For BM25 retrieval
+from rank_bm25 import BM25Okapi
 
 nltk.download('punkt')
 
-# Use a strong RoBERTa model
-model_name = "deepset/roberta-large-squad2"  # More powerful than base
+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)
@@ -31,7 +30,7 @@ def retrieve_relevant_text_bm25(question, sentences, top_n=3):
         bm25 = BM25Okapi(tokenized_corpus)
         tokenized_query = question.split()
         doc_scores = bm25.get_scores(tokenized_query)
-        top_n_indices = np.argsort(doc_scores)[::-1][:top_n]  # Get indices of top N
+        top_n_indices = np.argsort(doc_scores)[::-1][:top_n]
         relevant_texts = [sentences[i] for i in top_n_indices]
         return " ".join(relevant_texts)
     except Exception as e:
@@ -44,42 +43,41 @@ def answer_question(pdf, question, num_words):
             return text
 
         sentences = sent_tokenize(text)
-        relevant_text = retrieve_relevant_text_bm25(question, sentences) # Use BM25
+        relevant_text = retrieve_relevant_text_bm25(question, sentences)
 
         response = qa_pipeline(question=question, context=relevant_text)
         answer = response['answer']
 
         words = answer.split()
-        if len(words) > num_words:
-            answer = " ".join(words[:num_words])
-        elif len(words) < num_words:
+        if len(words) >= num_words:
+            answer = " ".join(words[:num_words])  # Strict truncation
+        elif relevant_text: #only add context if there is relevant text
             remaining_words = num_words - len(words)
-            added_sentences = []
+            added_words = 0
             for sentence in sent_tokenize(relevant_text):
-                if remaining_words > 0:
-                    sentence_words = sentence.split()
-                    to_add = min(remaining_words, len(sentence_words))
-                    added_sentences.append(" ".join(sentence_words[:to_add]))
-                    remaining_words -= to_add
-                else:
+                sentence_words = sentence.split()
+                words_to_add = min(remaining_words - added_words, len(sentence_words))
+                words.extend(sentence_words[:words_to_add])
+                added_words += words_to_add
+                if added_words == remaining_words:
                     break
-            answer += " " + " ".join(added_sentences)
-        answer = answer.strip()
-        return answer
+            answer = " ".join(words)
+            if len(answer.split()) > num_words: #truncate again if needed
+                answer = " ".join(answer.split()[:num_words])
+        return answer.strip()
 
     except Exception as e:
         return str(e)
 
-iface = gr.Interface(
-    fn=answer_question,
-    inputs=[
-        gr.File(type="filepath", label="Upload PDF"),
-        gr.Textbox(lines=2, placeholder="Ask a question"),
-        gr.Slider(minimum=1, maximum=1000, value=100, step=1, label="Number of Words")
-    ],
-    outputs=gr.Textbox(label="Answer"),
-    title="PDF Q&A with RoBERTa | Made by NP"
-)
+with gr.Blocks() as iface: #use blocks for custom layout
+    gr.Markdown("PDF Q&A with RoBERTa | Made by NP")
+    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
+    btn = gr.Button("Submit")
+    btn.click(fn=answer_question, inputs=[pdf_input, question_input, num_words_slider], outputs=answer_output)
 
 if __name__ == "__main__":
     iface.launch()