Update app.py

app.py

@@ -1,36 +1,37 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-import string
-import numpy as np
 from transformers import AutoTokenizer
 import onnxruntime as ort
+import numpy as np
+import string
+from huggingface_hub import InferenceClient
 import os
 
-# Initialize client and models
 client = InferenceClient(api_key=os.environ.get('HF_TOKEN'))
 
-# Constants for EOU calculation
-PUNCS = string.punctuation.replace("'", "")
-MAX_HISTORY = 4
-MAX_HISTORY_TOKENS = 1024
-EOU_THRESHOLD = 0.5
-
-# Initialize tokenizer and ONNX session
+# Model and ONNX setup
 HG_MODEL = "livekit/turn-detector"
 ONNX_FILENAME = "model_quantized.onnx"
+PUNCS = string.punctuation.replace("'", "")
+MAX_HISTORY = 4  # Adjusted to use the last 4 messages
+MAX_HISTORY_TOKENS = 512
+EOU_THRESHOLD = 0.5  # Updated threshold to match original
+
+# Initialize ONNX model
 tokenizer = AutoTokenizer.from_pretrained(HG_MODEL)
 onnx_session = ort.InferenceSession(ONNX_FILENAME, providers=["CPUExecutionProvider"])
 
-# Helper functions for EOU
+# Softmax function
 def softmax(logits):
     exp_logits = np.exp(logits - np.max(logits))
     return exp_logits / np.sum(exp_logits)
 
+# Normalize text
 def normalize_text(text):
     def strip_puncs(text):
         return text.translate(str.maketrans("", "", PUNCS))
     return " ".join(strip_puncs(text).lower().split())
 
+# Format chat context
 def format_chat_ctx(chat_ctx):
     new_chat_ctx = []
     for msg in chat_ctx:
@@ -39,14 +40,19 @@ def format_chat_ctx(chat_ctx):
             if content:
                 msg["content"] = content
                 new_chat_ctx.append(msg)
+
+    # Tokenize with chat template
     convo_text = tokenizer.apply_chat_template(
         new_chat_ctx, add_generation_prompt=False, add_special_tokens=False, tokenize=False
     )
+
+    # Remove EOU token from the current utterance
     ix = convo_text.rfind("<|im_end|>")
     return convo_text[:ix]
 
+# Calculate EOU probability
 def calculate_eou(chat_ctx, session):
-    formatted_text = format_chat_ctx(chat_ctx[-MAX_HISTORY:])
+    formatted_text = format_chat_ctx(chat_ctx[-MAX_HISTORY:])  # Use the last 4 messages
     inputs = tokenizer(
         formatted_text,
         return_tensors="np",
@@ -61,35 +67,42 @@ def calculate_eou(chat_ctx, session):
     return probs[eou_token_id]
 
 
+# Respond function
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    max_tokens,
+    temperature,
+    top_p,
+):
+    # Keep the last 4 conversation pairs (user-assistant)
+    messages = [{"role": "system", "content": os.environ.get("CHARACTER_DESC")}]
 
-messages = []
-
-
-def chatbot(user_input):
-    global messages
-    
-    # Exit condition
-    if user_input.lower() == "exit":
-        messages = []  # Reset conversation history
-        return "Chat ended. Refresh the page to start again."
+    for val in history[-20:]:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
 
-    
-    # Add user message to conversation history
-    messages.append({"role": "user", "content": user_input})
+    # Add the new user message to the context
+    messages.append({"role": "user", "content": message})
 
-    # Calculate EOU to determine if user has finished typing
+    # Calculate EOU probability
     eou_prob = calculate_eou(messages, onnx_session)
+    print(f"EOU Probability: {eou_prob}")  # Debug output
+
+    # If EOU is below the threshold, ask for more input
     if eou_prob < EOU_THRESHOLD:
-        yield "[I'm waiting for you to complete the sentence...]"
+        yield "[Waiting for user to continue input...]"
         return
 
-    # Stream the chatbot's response
+
     stream = client.chat.completions.create(
-        model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
+        model=os.environ.get('MODEL_ID'),
         messages=messages,
-        temperature=0.6,
-        max_tokens=2200,
-        top_p=0.95,
+        temperature = 0.6,
+        max_tokens= 2048,
+        top_p = 0.9,
         stream=True
     )
 
@@ -97,23 +110,23 @@ def chatbot(user_input):
     for chunk in stream:
         bot_response += chunk.choices[0].delta.content
         yield bot_response
+    
 
-    # Add final bot response to conversation history
-    messages.append({"role": "assistant", "content": bot_response})
-
-# Create Gradio interface
-with gr.Blocks(theme='darkdefault') as demo:
-    gr.Markdown("""# Chat with DeepSeek""")
-
-    with gr.Row():
-        with gr.Column():
-            user_input = gr.Textbox(label="Your Message", placeholder="Type your message here...")
-            submit_button = gr.Button("Send")
-        with gr.Column():
-            chat_output = gr.Textbox(label="Chatbot Response", interactive=False)
-
-    # Define interactions
-    submit_button.click(chatbot, inputs=[user_input], outputs=[chat_output])
 
-# Launch the app
-demo.launch()
+# Gradio interface
+demo = gr.ChatInterface(
+    respond,
+    # additional_inputs=[
+    #     # Commented out to disable user modification of the system message
+    #     # gr.Textbox(value="You are an assistant.", label="System message"),
+    #     gr.Slider(minimum=1, maximum=4096, value=256, step=1, label="Max new tokens"),
+    #     gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+    #     gr.Slider(
+    #         minimum=0.1,
+    #         maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"
+    #     ),
+    # ],
+)
+
+if __name__ == "__main__":
+    demo.launch()