Update app.py

app.py

@@ -1,8 +1,9 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import fitz  # PyMuPDF
+import torch
 
-# Define a function to load the PDF document and cache the content with Streamlit's caching
+# Function to load the PDF document
 @st.cache(allow_output_mutation=True)
 def load_pdf_document(file_path):
     text = ""
@@ -11,40 +12,44 @@ def load_pdf_document(file_path):
             text += page.get_text()
     return text
 
-# Define a function to get answers from the model and cache the model and tokenizer
+# Function to load the model and tokenizer
 @st.cache(allow_output_mutation=True)
 def load_model_and_tokenizer(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
     return tokenizer, model
 
-# UI for the Streamlit app
-st.title("Question Answering with LLaMA 2")
+# Function to generate an answer from the model
+def generate_answer(context, query, tokenizer, model):
+    # Preprocess and truncate the context to fit within model limits
+    encoded_input = tokenizer.encode_plus(query, context, add_special_tokens=True, return_tensors="pt", truncation=True, max_length=tokenizer.model_max_length - 20)
+    input_ids = encoded_input["input_ids"]
+    attention_mask = encoded_input["attention_mask"]
+    
+    # Generate an answer using max_new_tokens to limit output length
+    output = model.generate(input_ids, attention_mask=attention_mask, max_new_tokens=150, num_return_sequences=1, temperature=0.7, top_p=0.9)
+    answer = tokenizer.decode(output[0], skip_special_tokens=True)
+    return answer
 
-# Load the document
-document_path = "jeff_wo.pdf"  # Change this to the path of your PDF document in the repository
+# Streamlit UI
+st.title("Question Answering with LLaMA 2")
+document_path = "jeff_wo.pdf"
 document_text = load_pdf_document(document_path)
-# Displaying the document text can be optional, based on your preference or usability considerations
-st.text_area("Document Text", value=document_text, height=300, help="Content of the PDF document.")
+
+# Optional: Display the document text or a portion of it
+st.text_area("Document Text (preview)", value=document_text[:1000], height=250, help="Preview of the document text.")
 
 # Load model and tokenizer
 model_name = "NousResearch/Llama-2-7b-chat-hf"
 tokenizer, model = load_model_and_tokenizer(model_name)
 
-# Sidebar for user input
-st.sidebar.header("Ask a Question")
-query = st.sidebar.text_input("Enter your question:", "")
+# User input for the query
+query = st.text_input("Enter your question:", "")
 
-if st.sidebar.button("Answer"):
+if st.button("Generate Answer"):
     if query:
         with st.spinner("Generating answer..."):
-            input_text = f"Context: {document_text}\nQ: {query}\nA:"
-            input_ids = tokenizer.encode(input_text, return_tensors="pt")
-            # Adjust the generation parameters as needed
-            output = model.generate(input_ids, max_length=512, num_return_sequences=1, temperature=0.7, top_p=0.9)
-            answer = tokenizer.decode(output[0], skip_special_tokens=True)
+            answer = generate_answer(document_text, query, tokenizer, model)
             st.write(answer)
     else:
-        st.sidebar.error("Please enter a question.")
-else:
-    st.write("Enter a question on the sidebar and click 'Answer' to get a response.")
+        st.error("Please enter a question to get an answer.")