Update app.py

app.py

@@ -19,14 +19,24 @@ tokenizer = AutoTokenizer.from_pretrained(model_dir)
 #     offload_folder=offload_dir # Directory to offload layers (for larger models).
 # )
 
+## CPU
+# base_model = AutoModelForCausalLM.from_pretrained(
+#     model_dir,
+#     device_map="cpu",  # Force CPU usage
+#     torch_dtype=torch.float32,  # Use float32 for CPU
+#     trust_remote_code=True,
+#     offload_folder=offload_dir # Directory to offload layers (for larger models).
+# )
+
+## GPU
 base_model = AutoModelForCausalLM.from_pretrained(
     model_dir,
-    device_map="cpu",  # Force CPU usage
-    torch_dtype=torch.float32,  # Use float32 for CPU
-    trust_remote_code=True,
-    offload_folder=offload_dir # Directory to offload layers (for larger models).
+    device_map="auto",         # Automatically uses available GPU(s)
+    torch_dtype=torch.float16, # Use half precision for faster inference on GPU
+    trust_remote_code=True
 )
 
+
 # Load the adapter (PEFT) weights.
 model = PeftModel.from_pretrained(base_model, model_dir)