Update app.py

app.py

@@ -3,18 +3,6 @@ from PyPDF2 import PdfReader
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Initialize the tokenizer and model from the saved checkpoint
-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
-)
-
-# Use GPU if available
-if torch.cuda.is_available():
-    model.to("cuda")
-
 # Set up the Streamlit app layout
 st.set_page_config(page_title="RAG PDF Chatbot", layout="wide")
 
@@ -58,6 +46,14 @@ submit_button = st.button("Submit")
 if 'chat_history' not in st.session_state:
     st.session_state.chat_history = []
 
+# Load the tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("himmeow/vi-gemma-2b-RAG")
+model = AutoModelForCausalLM.from_pretrained("himmeow/vi-gemma-2b-RAG")
+
+# Use GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+
 # Function to extract text from PDF files
 def extract_text_from_pdfs(files):
     text = ""
@@ -81,15 +77,11 @@ if submit_button and query:
         """
 
         # Encode the input text
-        input_ids = tokenizer(prompt, return_tensors="pt")
-
-        # Use GPU for input ids if available
-        if torch.cuda.is_available():
-            input_ids = input_ids.to("cuda")
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
 
         # Generate the response
         outputs = model.generate(
-            **input_ids,
+            input_ids=input_ids,
             max_new_tokens=500,
             no_repeat_ngram_size=5,
         )