Update app.py

app.py

@@ -2,27 +2,41 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Load model and tokenizer from Hugging Face Hub
+# Set model ID
 model_id = "deepseek-ai/deepseek-coder-1.3b-base"
-# model_id = "gpt2"
+# For smaller testing you can use: model_id = "gpt2"
+
+# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+)
+
+# Move model to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
 def generate_code(prompt):
     if not prompt.strip():
         return "⚠ Please enter a valid prompt."
-    
-    inputs = tokenizer(prompt, return_tensors="pt")
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
-    
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
     with torch.no_grad():
-        outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7)
-    
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.7
+        )
+
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-demo = gr.Interface(fn=generate_code, 
-                    inputs=gr.Textbox(lines=5, label="Enter Prompt"), 
-                    outputs="text", 
-                    title="Code Generator using DeepSeek")
+demo = gr.Interface(
+    fn=generate_code,
+    inputs=gr.Textbox(lines=5, label="Enter Prompt"),
+    outputs=gr.Textbox(label="Generated Output"),
+    title="Code Generator using DeepSeek"
+)
 
 demo.launch()