Update app.py

app.py

@@ -436,6 +436,8 @@ For more information on `huggingface_hub` Inference API support, please check th
 
 # app.py
 
+# app.py
+
 import os
 import gradio as gr
 from huggingface_hub import InferenceClient
@@ -447,7 +449,7 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 
 # Initialize Hugging Face Inference Client
 client = InferenceClient(
-    model="mistralai/Mamba-Codestral-7B-v0.1",
+    model="mistralai/Codestral-22B-v0.1",
     token=HF_TOKEN
 )
 
@@ -459,36 +461,32 @@ system_message = (
     "based on their requirements."
 )
 
-# Template for conversation history
-def format_prompt(message, history):
-    prompt = f"System: {system_message}\n"
-    for user_msg, bot_msg in history:
-        prompt += f"User: {user_msg}\n"
-        if bot_msg:
-            prompt += f"Assistant: {bot_msg}\n"
-    prompt += f"User: {message}\nAssistant:"
-    return prompt
-
-# Streaming chatbot logic using text generation
+# Streaming chatbot logic using chat.completions
 def respond(message, history):
-    prompt = format_prompt(message, history)
+    # Prepare messages with system prompt
+    messages = [{"role": "system", "content": system_message}]
+    for msg in history:
+        messages.append(msg)
+    messages.append({"role": "user", "content": message})
 
+    # Stream response from the model
     response = ""
-    for chunk in client.text_generation(
-        prompt=prompt,
-        max_new_tokens=1024,
+    for chunk in client.chat.completions.create(
+        model="mistralai/Codestral-22B-v0.1",
+        messages=messages,
+        max_tokens=1024,
         temperature=0.7,
         top_p=0.95,
-        stop_sequences=["\nUser:", "\nSystem:"],
-        stream=True
+        stream=True,
     ):
-        response += chunk
+        token = chunk.choices[0].delta.get("content", "") or ""
+        response += token
         yield response
 
 # Create Gradio interface
 with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(type='messages')  # Still use modern message format
-    gr.ChatInterface(fn=respond, chatbot=chatbot, type="messages")
+    chatbot = gr.Chatbot(type='messages')  # Use modern message format
+    gr.ChatInterface(fn=respond, chatbot=chatbot, type="messages")  # Match format
 
 # Launch app
 if __name__ == "__main__":