check if MPS available

app.py

@@ -3,6 +3,10 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
+
+print("CUDA available: ", torch.cuda.is_available())
+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(
@@ -12,8 +16,8 @@ model = AutoModelForCausalLM.from_pretrained(
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 
-# Use CPU if CUDA is not available
-device = torch.device("cpu" if not torch.cuda.is_available() else "cuda")
+# 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)
 
 # Theme builder
@@ -39,7 +43,7 @@ def chatbot_response(query):
     inputs = tokenizer.apply_chat_template(
         messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
 
-    outputs = model.generate(inputs, max_new_tokens=512, do_sample=True, top_k=50,
+    outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, top_k=50,
                              top_p=0.95, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
     model_response = tokenizer.decode(
         outputs[0][len(inputs[0]):], skip_special_tokens=True)

requirements.txt

@@ -3,4 +3,6 @@ gradio==4.44.1
 requests
 transformers
 minijinja
-torch --extra-index-url https://download.pytorch.org/whl/cu118
+torch --extra-index-url https://download.pytorch.org/whl/torch-cuda80
+torchvision
+torchaudio