Update app.py

app.py

@@ -4,28 +4,23 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 model_id = "deepseek-ai/DeepSeek-V3"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True) # device_map="auto" for GPU if available, ADD trust_remote_code=True
-
-print("Model and tokenizer loaded successfully!")
-
-
-def generate_text(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True) # ADD trust_remote_code=True
+
+def predict(message, history):
+    prompt = tokenizer.apply_chat_template(
+        [{"role": "user", "content": message}],
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda") # Assuming you have CUDA available in your Space
     outputs = model.generate(**inputs, max_new_tokens=50) # Adjust max_new_tokens as needed
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Example usage (for testing in your app.py):
-if __name__ == "__main__":
-    prompt = "Write a short story about a cat who can talk."
-    response = generate_text(prompt)
-    print(f"Prompt: {prompt}\nResponse: {response}")
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
 iface = gr.ChatInterface(
     fn=predict,
-    inputs=gr.Chatbox(label="Chat with DeepSeek-V3"),
-    outputs=gr.Chatbot(label="DeepSeek-V3"),
-    title="DeepSeek-V3 Chatbot",
-    description="Chat with the DeepSeek-V3 model from Hugging Face.",
+    inputs=gr.Chatbox(),
+    outputs=gr.Chatbot()
 )
 
 iface.launch()