Update app.py

app.py

@@ -1,41 +1,49 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import InferenceClient
 
-# Load the model and tokenizer
-model_name = "deepseek-ai/DeepSeek-R1"
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+# Initialize the InferenceClient with the model hosted on Hugging Face Hub
+client = InferenceClient("deepseek-ai/DeepSeek-R1")
 
 def respond(message, history: list[tuple[str, str]]):
+    # System message and generation parameters
+    system_message = "You are a friendly Chatbot."
+    max_tokens = 2048
+    temperature = 0.7
+    top_p = 0.95
+
     # Prepare the conversation history
-    messages = []
+    messages = [{"role": "system", "content": system_message}]
+
     for user_msg, assistant_msg in history:
         if user_msg:
             messages.append({"role": "user", "content": user_msg})
         if assistant_msg:
             messages.append({"role": "assistant", "content": assistant_msg})
-    messages.append({"role": "user", "content": message})
-
-    # Tokenize the input
-    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt")
 
-    # Generate the response
-    outputs = model.generate(inputs, max_length=2048, temperature=0.7, top_p=0.95, do_sample=True)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    return response
-
-# Custom ChatInterface with undo and retry buttons
-def chat_interface(message, history):
-    return respond(message, history)
+    # Add the current user message
+    messages.append({"role": "user", "content": message})
 
-# Create the Gradio interface
+    # Stream the response from the model
+    response = ""
+    for message in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        token = message.choices[0].delta.content
+        response += token
+        yield response
+
+# Create the Gradio ChatInterface with Retry, Undo, and Clear buttons
 demo = gr.ChatInterface(
-    fn=chat_interface,
-    retry_btn="Retry",
-    undo_btn="Undo",
-    clear_btn="Clear",
+    fn=respond,  # Function to handle chat responses
+    retry_btn="Retry",  # Add a Retry button
+    undo_btn="Undo",  # Add an Undo button
+    clear_btn="Clear",  # Add a Clear button
 )
 
+# Launch the Gradio app
 if __name__ == "__main__":
     demo.launch()