Update app.py

app.py

@@ -21,8 +21,18 @@ vectordb = FAISS.from_documents(split_docs, embeddings)
 
 # Load model and create pipeline
 model_name = "01-ai/Yi-Coder-9B-Chat"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto")
+#model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cuda", torch_dtype="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    device_map=device,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+)
+#model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto")
 qa_pipeline = pipeline(
     "text-generation",
     model=model,