Update app.py

app.py

@@ -6,41 +6,33 @@ import string
 from huggingface_hub import InferenceClient
 import os
 
-# Initialize client with error handling
-try:
-    qwen_client = InferenceClient(
-        "nztinversive/llama3.2-1b-Uncensored",
-        timeout=120
-    )
-except Exception as e:
-    print(f"Error initializing client: {e}")
-    raise
+# Initialize Qwen client
+qwen_client = InferenceClient("huihui-ai/SmolLM2-1.7B-Instruct-abliterated")
 
 # Model and ONNX setup
 HG_MODEL = "livekit/turn-detector"
 ONNX_FILENAME = "model_quantized.onnx"
 PUNCS = string.punctuation.replace("'", "")
-MAX_HISTORY = 4
+MAX_HISTORY = 4  # Adjusted to use the last 4 messages
 MAX_HISTORY_TOKENS = 512
-EOU_THRESHOLD = 0.5
+EOU_THRESHOLD = 0.5  # Updated threshold to match original
 
 # Initialize ONNX model
-try:
-    tokenizer = AutoTokenizer.from_pretrained(HG_MODEL)
-    onnx_session = ort.InferenceSession(ONNX_FILENAME, providers=["CPUExecutionProvider"])
-except Exception as e:
-    print(f"Error initializing models: {e}")
-    raise
+tokenizer = AutoTokenizer.from_pretrained(HG_MODEL)
+onnx_session = ort.InferenceSession(ONNX_FILENAME, providers=["CPUExecutionProvider"])
 
+# 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:
@@ -50,15 +42,18 @@ def format_chat_ctx(chat_ctx):
                 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",
@@ -72,71 +67,67 @@ def calculate_eou(chat_ctx, session):
     eou_token_id = tokenizer.encode("<|im_end|>")[-1]
     return probs[eou_token_id]
 
-def respond(message, history, max_tokens, temperature, top_p):
-    messages = [{"role": "system", "content": os.environ.get("CHARACTER_DESC", "You are a helpful assistant.")}]
-    
-    # Process history
-    for user_msg, assistant_msg in history[-10:]:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
-    
-    # Add new message
+
+# 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")}]
+
+    for val in history[-10:]:  # Only use the last 4 pairs
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
+
+    # Add the new user message to the context
     messages.append({"role": "user", "content": message})
 
     # Calculate EOU probability
-    try:
-        eou_prob = calculate_eou(messages, onnx_session)
-        print(f"EOU Probability: {eou_prob}")
-        
-        if eou_prob < EOU_THRESHOLD:
-            yield "[Waiting for additional input...]"
-            return
-    except Exception as e:
-        print(f"EOU calculation error: {e}")
-        yield "[Error in conversation analysis]"
+    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 "[Waiting for user to continue input...]"
         return
 
-    # Generate response
-    try:
-        # Format prompt for text generation
-        prompt = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )
-        
-        stream = qwen_client.text_generation(
-            prompt,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            stream=True,
-            details=True,
-        )
-
-        response = ""
-        for chunk in stream:
-            if chunk.token.text:
-                response += chunk.token.text
-                yield response
-
-        print(f"Final response: {response}")
-        
-    except Exception as e:
-        print(f"Generation error: {e}")
-        yield "[Error generating response]"
+    # Generate response with Qwen
+    response = ""
+    for message in qwen_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
+    
+    print(f"Generated response: {response}")
+
 
 # Gradio interface
 demo = gr.ChatInterface(
     respond,
-    additional_inputs=[
-        gr.Slider(1, 4096, value=256, label="Max Tokens"),
-        gr.Slider(0.1, 4.0, value=0.7, label="Temperature"),
-        gr.Slider(0.1, 1.0, value=0.95, label="Top-p"),
-    ]
+    # 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()
+    demo.launch()
+