Update app.py

app.py

@@ -6,44 +6,34 @@ st.set_page_config(page_title="Question Answering Tool", page_icon=":mag_right:"
 
 @st.cache_resource
 def load_model():
+    """Loads the DistilBERT model and tokenizer for QA."""
     model = DistilBertForQuestionAnswering.from_pretrained("distilbert-base-uncased-distilled-squad")
     tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased-distilled-squad")
     return model, tokenizer
 
-def generate_prompt(question, text):
-    """Enhance the input prompt to guide the model better."""
-    return (
-        f"Context: {text}\n\n"
-        f"Instruction: Read the above context carefully and extract the most relevant answer.\n"
-        f"Question: {question}\n"
-        f"Answer:"
-    )
-
 def get_answer(question, text, tokenizer, model):
-    # Special case for bot identity
+    """Extracts the most relevant answer from the given text."""
     if any(phrase in question.lower() for phrase in ["your name", "who are you", "about you"]):
         return "I am Numini, NativUttarMini, created by Sanju Debnath at University of Calcutta."
 
-    # Enhance the input for better response
-    prompt_text = generate_prompt(question, text)
-
-    # Tokenize with truncation for better handling of large text
-    inputs = tokenizer(prompt_text, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    # Tokenize input text and question
+    inputs = tokenizer(question, text, return_tensors="pt", truncation=True, padding=True, max_length=512)
 
     with torch.no_grad():
         outputs = model(**inputs)
 
-    start = torch.argmax(outputs.start_logits)
-    end = torch.argmax(outputs.end_logits) + 1
+    start_idx = torch.argmax(outputs.start_logits)
+    end_idx = torch.argmax(outputs.end_logits) + 1
 
-    # Handle cases where model is uncertain
-    if start >= end:
+    # Validate extracted indices
+    if start_idx >= end_idx or end_idx > inputs.input_ids.shape[1]:
         return "I couldn't find a clear answer in the given text."
 
-    answer = tokenizer.decode(inputs.input_ids[0][start:end], skip_special_tokens=True)
+    # Decode extracted answer
+    answer = tokenizer.decode(inputs.input_ids[0][start_idx:end_idx], skip_special_tokens=True)
 
-    # Post-process answer for better readability
-    if len(answer.split()) < 3:  # Model sometimes returns incomplete answers
+    # Ensure answer is meaningful
+    if len(answer.split()) < 2:
         return "I'm not sure about the exact answer. Can you try rephrasing the question?"
     
     return answer
@@ -58,15 +48,17 @@ def main():
         text = st.text_area("📜 Enter the text/document:", height=200)
         question = st.text_input("❓ Enter your question:")
 
-        if st.form_submit_button("🔍 Get Answer"):
-            if not text.strip():
-                st.warning("⚠️ Please enter some text to analyze.")
-            elif not question.strip():
-                st.warning("⚠️ Please enter a question.")
-            else:
-                st.text("🤖 Thinking...")
+        submit = st.form_submit_button("🔍 Get Answer")
+    
+    if submit:
+        if not text.strip():
+            st.warning("⚠️ Please enter some text to analyze.")
+        elif not question.strip():
+            st.warning("⚠️ Please enter a question.")
+        else:
+            with st.spinner("🤖 Thinking..."):
                 answer = get_answer(question, text, tokenizer, model)
-                st.success(f"✅ Answer: {answer}")
-
-main()
+            st.success(f"✅ Answer: {answer}")
 
+if __name__ == "__main__":
+    main()