Update app.py

app.py

@@ -1,10 +1,18 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import pandas as pd
 import re
 
-# Load the Question Answering model
-qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
+# Load the LLaMA model and tokenizer
+@st.cache_resource
+def load_llama_model():
+    """Load the LLaMA model and tokenizer."""
+    model_name = "meta-llama/Llama-2-7b-chat-hf"  # Replace with your preferred LLaMA model
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+qa_pipeline = load_llama_model()
 
 # Load SOP Dataset
 @st.cache_data
@@ -13,7 +21,6 @@ def load_sop_dataset():
     dataset = pd.read_csv("dataset.csv")  # Ensure this file is uploaded to your Hugging Face Space
     return dataset
 
-# Load the dataset
 dataset = load_sop_dataset()
 
 # Utility function to find the most relevant context
@@ -32,7 +39,7 @@ def find_best_context(question, dataset):
     return best_context
 
 # Streamlit UI
-st.title("SOP Question Answering AI")
+st.title("SOP Question Answering AI with LLaMA")
 st.markdown("Ask any question about Standard Operating Procedures:")
 
 # User input
@@ -47,10 +54,10 @@ if st.button("Get Answer"):
 
             if context:
                 with st.spinner("Answering your question..."):
-                    result = qa_pipeline(question=question, context=context)
+                    prompt = f"Context: {context}\n\nQuestion: {question}\nAnswer:"
+                    result = qa_pipeline(prompt, max_length=150, num_return_sequences=1)
                     st.success("Answer:")
-                    st.write(result["answer"])
-                    st.write("Confidence Score:", result["score"])
+                    st.write(result[0]["generated_text"])
             else:
                 st.warning("No relevant context found. Please try rephrasing your question.")
     else: