updates

app.py

@@ -38,16 +38,25 @@ print("MPS available: ", torch.backends.mps.is_available())
 tokenizer = AutoTokenizer.from_pretrained(
     "deepseek-ai/deepseek-coder-1.3b-instruct", trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    "deepseek-ai/deepseek-coder-1.3b-instruct", trust_remote_code=True, torch_dtype=torch.bfloat16)
+    "deepseek-ai/deepseek-coder-1.3b-instruct", 
+    trust_remote_code=True, 
+    torch_dtype=torch.float16,  # Use float16 for better GPU memory efficiency
+    device_map="auto"  # Automatically handle device placement
+)
 
 # Disable tokenizers parallelism warning
 os.environ["TOKENIZERS_PARALLELISM"] = "True"
 
 
-# Check if MPS (Metal Performance Shaders) is available
-device = torch.device(
-    "mps") if torch.backends.mps.is_available() else torch.device("cpu")
-model = model.to(device)
+# Configure device
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    # Print GPU information
+    print(f"Using GPU: {torch.cuda.get_device_name(0)}")
+    print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
+else:
+    device = torch.device("cpu")
+    print("No GPU available, using CPU")
 
 
 # Function to handle user input and generate a response

requirements.txt

@@ -6,3 +6,4 @@ minijinja
 torch --extra-index-url https://download.pytorch.org/whl/cu118
 torchvision
 torchaudio
+accelerate==0.26.0