Update app.py

app.py

@@ -6,25 +6,33 @@ import fitz  # PyMuPDF
 # Load the tokenizer and model
 @st.cache_resource
 def load_model():
-    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
+    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, model = load_model()
 
 # Function to read text from a PDF file
 def read_pdf(file):
-    text = ""
-    with fitz.open("pdf", file.read()) as doc:
-        for page in doc:
-            text += page.get_text()
-    return text
+    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
 
 # Streamlit app
 st.title("PDF Question Answering with vi-gemma-2b-RAG")
@@ -37,27 +45,34 @@ if uploaded_file is not None and question:
     # Read PDF content
     pdf_text = read_pdf(uploaded_file)
     
-    # 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)
-
-    # Generate a response
-    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)
+    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)
+
+        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.")