Update app.py

app.py

@@ -15,15 +15,15 @@ def extract_text_from_pdf(uploaded_file):
         pdf_text += page.extract_text()
     return pdf_text
 
-# Initialize the tokenizer and model
+# Initialize the tokenizer and model on CPU first
 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
 )
 
-# Use GPU if available
+# Move model to GPU if available
 if torch.cuda.is_available():
     model.to("cuda")