Update app.py

app.py

@@ -1,78 +1,71 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-import fitz  # PyMuPDF
+import fitz  # PyMuPDF for PDF handling
 
-# Load the tokenizer and model
+# Load the model and tokenizer
 @st.cache_resource
 def load_model():
-    try:
-        tokenizer = AutoTokenizer.from_pretrained("ricepaper/vi-gemma-2b-RAG")
-        model = AutoModelForCausalLM.from_pretrained(
-            "ricepaper/vi-gemma-2b-RAG",
-            device_map="auto",
-            torch_dtype=torch.bfloat16
-        )
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        model.to(device)
-        return tokenizer, model
-    except Exception as e:
-        st.error(f"Error loading model: {e}")
-        return None, None
+    tokenizer = AutoTokenizer.from_pretrained("himmeow/vi-gemma-2b-RAG")
+    model = AutoModelForCausalLM.from_pretrained(
+        "himmeow/vi-gemma-2b-RAG",
+        device_map="auto",
+        torch_dtype=torch.bfloat16
+    )
+    if torch.cuda.is_available():
+        model.to("cuda")
+    return tokenizer, model
 
-tokenizer, model = load_model()
+# Function to extract text from PDF
+def extract_text_from_pdf(pdf_file):
+    doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
+    text = ""
+    for page_num in range(doc.page_count):
+        page = doc.load_page(page_num)
+        text += page.get_text()
+    return text
 
-# Function to read text from a PDF file
-def read_pdf(file):
-    try:
-        text = ""
-        with fitz.open("pdf", file.read()) as doc:
-            for page in doc:
-                text += page.get_text()
-        return text
-    except Exception as e:
-        st.error(f"Error reading PDF file: {e}")
-        return None
+# Function to generate response from model
+def generate_response(input_text, query, tokenizer, model):
+    prompt = """
+    ### Instruction and Input:
+    Based on the following context/document:
+    {}
+    Please answer the question: {}
 
-# Streamlit app
-st.title("PDF Question Answering with vi-gemma-2b-RAG")
-st.write("Upload a PDF file, and ask a question based on its content.")
-
-uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
-question = st.text_input("Enter your question:")
+    ### Response:
+    {}
+    """
+    formatted_input = prompt.format(input_text, query, " ")
+    input_ids = tokenizer(formatted_input, return_tensors="pt")
+    if torch.cuda.is_available():
+        input_ids = input_ids.to("cuda")
+    outputs = model.generate(
+        **input_ids,
+        max_new_tokens=500,
+        no_repeat_ngram_size=5
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-if uploaded_file is not None and question:
-    # Read PDF content
-    pdf_text = read_pdf(uploaded_file)
+# Streamlit app
+def main():
+    st.title("PDF Question Answering with vi-gemma-2b-RAG")
     
-    if pdf_text:
-        # Prepare the input for the model
-        prompt_template = """
-        ### Instruction and Input:
-        Based on the following context/documentation:
-        {}
-        Please answer the question: {}
-
-        ### Response:
-        {}
-        """
-        input_text = prompt_template.format(pdf_text, question, "")
-        input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)
+    pdf_file = st.file_uploader("Upload a PDF file", type=["pdf"])
+    
+    if pdf_file is not None:
+        with st.spinner("Reading the PDF..."):
+            pdf_text = extract_text_from_pdf(pdf_file)
+        
+        st.text_area("Extracted Text", pdf_text, height=300)
+        
+        query = st.text_input("Enter your question:")
+        
+        if st.button("Get Answer"):
+            with st.spinner("Generating response..."):
+                tokenizer, model = load_model()
+                response = generate_response(pdf_text, query, tokenizer, model)
+                st.text_area("Response", response, height=200)
 
-        try:
-            # Generate a response
-            with torch.no_grad():  # Disable gradient calculation for inference
-                with torch.cuda.amp.autocast():
-                    outputs = model.generate(
-                        **input_ids,
-                        max_new_tokens=200,
-                        no_repeat_ngram_size=5
-                    )
-            
-            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            st.subheader("Answer:")
-            st.write(response)
-        except Exception as e:
-            st.error(f"Error generating response: {e}")
-    else:
-        st.error("Unable to read text from the uploaded PDF file.")
+if __name__ == "__main__":
+    main()