Update app.py

app.py

@@ -1,30 +1,41 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import fitz  # PyMuPDF
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Streamlit's cache decorator to cache the model and tokenizer loading
+# Function to read and extract text from a PDF document
+def read_pdf(file_path):
+    text = ""
+    with fitz.open(file_path) as doc:
+        for page in doc:
+            text += page.get_text()
+    return text
 
-def load_pipeline():
-    model_name = "NousResearch/Llama-2-7b-chat-hf"  # Replace with your actual model name
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    chat_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=300)
-    return chat_pipeline
+# Load the document text
+document_text = read_pdf("jeff_wo.pdf")  # Adjust path to your PDF file
 
-# Initialize the pipeline
-chat_pipeline = load_pipeline()
+# Streamlit UI
+st.title("LLaMA 2-based Q&A System")
+st.write("### Enter your query below:")
+query = st.text_input("Query")
 
-st.title("Interact with Your Model")
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-chat-hf")
+model = AutoModelForCausalLM.from_pretrained("NousResearch/Llama-2-7b-chat-hf")
 
-# User input
-user_input = st.text_area("Enter your prompt:", "")
+# Function to get answers using the LLaMA 2 model
+def get_answer(context, query):
+    input_text = f"Context: {context}\nQ: {query}\nA:"
+    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    # Generate an answer to the query
+    output = model.generate(input_ids, max_length=512, num_return_sequences=1)
+    answer = tokenizer.decode(output[0], skip_special_tokens=True)
+    return answer
 
-if st.button("Submit"):
-    if user_input:
-        try:
-            # Generate text based on the input
-            generated_text = chat_pipeline(user_input)[0]['generated_text']
-            st.write(generated_text)
-        except Exception as e:
-            st.error(f"Error generating text: {e}")
-    else:
-        st.write("Please enter a prompt.")
+# Button to generate answers
+if st.button("Get Answer"):
+    with st.spinner("Finding the answer..."):
+        answer = get_answer(document_text, query)
+    st.write("### Answer:")
+    st.write(answer)
+
+# Ensure to replace `path/to/your/document.pdf` with the actual path to the PDF in your repo