Update app.py

app.py

@@ -1,29 +1,21 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import pipeline
 import pandas as pd
 import re
 
-# Load the Falcon-3B-Instruct model and tokenizer
+# Load the Question Answering model
 @st.cache_resource
-def load_falcon_model():
-    """Load the Falcon-3B-Instruct model and tokenizer."""
-    model_name = "tiiuae/falcon-3b-instruct"  # Smaller model for faster loading
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype="auto",  # Use FP16 if supported
-        device_map="auto"    # Automatically distributes across available devices
-    )
-    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+def load_qa_pipeline():
+    """Load the QA pipeline with deepset/roberta-base-squad2 model."""
+    return pipeline("question-answering", model="deepset/roberta-base-squad2")
 
-qa_pipeline = load_falcon_model()
+qa_pipeline = load_qa_pipeline()
 
 # Load SOP Dataset
 @st.cache_data
 def load_sop_dataset():
     """Load SOP dataset from CSV."""
-    dataset = pd.read_csv("dataset.csv")  # Ensure this file is uploaded to your Hugging Face Space
-    return dataset
+    return pd.read_csv("dataset.csv")  # Ensure this file is uploaded to your Hugging Face Space
 
 dataset = load_sop_dataset()
 
@@ -33,17 +25,17 @@ def find_best_context(question, dataset):
     best_score = 0
     best_context = None
 
-    for index, row in dataset.iterrows():
+    for _, row in dataset.iterrows():
         # Simple heuristic: Count the number of overlapping words
         overlap = len(set(question.lower().split()) & set(row["text"].lower().split()))
         if overlap > best_score:
             best_score = overlap
             best_context = row["text"]
-    
+
     return best_context
 
 # Streamlit UI
-st.title("SOP Question Answering AI with Falcon")
+st.title("SOP Question Answering AI")
 st.markdown("Ask any question about Standard Operating Procedures:")
 
 # User input
@@ -58,10 +50,10 @@ if st.button("Get Answer"):
 
             if context:
                 with st.spinner("Answering your question..."):
-                    prompt = f"Context: {context}\n\nQuestion: {question}\nAnswer:"
-                    result = qa_pipeline(prompt, max_length=150, num_return_sequences=1)
+                    result = qa_pipeline(question=question, context=context)
                     st.success("Answer:")
-                    st.write(result[0]["generated_text"])
+                    st.write(result["answer"])
+                    st.write("Confidence Score:", result["score"])
             else:
                 st.warning("No relevant context found. Please try rephrasing your question.")
     else: