Update app.py

app.py

@@ -1,38 +1,31 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 
-# Model and Tokenizer from Hugging Face Hub
-model_name = "deepseek-ai/DeepSeek-V3"
+model_id = "deepseek-ai/DeepSeek-V3"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto") # device_map="auto" for GPU if available
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.bfloat16,  # Use bfloat16 for faster and less memory-intensive inference if possible
-    trust_remote_code=True,      # Important for models with custom code
-    device_map="auto"           # Automatically use available GPU if possible
-)
+def predict(message, history):
+    conversation = []
+    for user_msg, bot_response in history:
+        conversation.append({"role": "user", "content": user_msg})
+        if bot_response: # Only add bot response if it exists
+            conversation.append({"role": "assistant", "content": bot_response})
+    conversation.append({"role": "user", "content": message})
 
-def generate_response(prompt, history=[]):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=500) # Adjust max_new_tokens as needed
+    inputs = tokenizer.apply_chat_template(conversation=conversation, tokenizer=tokenizer, return_tensors="pt").to("cuda" if model.device.type == 'cuda' else "cpu") # Move input to GPU if model is on GPU
+    outputs = model.generate(**inputs, max_new_tokens=512) # Adjust max_new_tokens as needed
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Basic chat history handling (optional, can be improved)
-    history.append((prompt, response)) # Append user prompt and model response to history
+    # Basic cleanup (you might need more sophisticated cleaning)
+    response = response.replace("<|assistant|>", "").strip()
     return response
 
 iface = gr.ChatInterface(
-    fn=generate_response,
-    inputs=gr.Chatbox(lines=7, placeholder="Type your message here..."),
-    outputs="text",
+    fn=predict,
+    inputs=gr.Chatbox(placeholder="Type a message..."),
+    outputs=gr.Chatbot(),
     title="DeepSeek-V3 Chatbot",
-    description="Chat with the DeepSeek-V3 model. Please be patient, initial loading might take a few minutes. For better performance, use a Space with a GPU.",
-    examples=[
-        "Hello, how are you?",
-        "What is the capital of France?",
-        "Tell me a joke."
-    ]
+    description="Chat with the DeepSeek-V3 model.",
 )
-
-iface.launch(share=False)
+iface.launch()