Update app.py

app.py

@@ -1,84 +1,90 @@
 import gradio as gr
-from PIL import Image
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
-import requests
+import numpy as np
+import string
 
-# Load the model and processor
-model = Qwen2VLForConditionalGeneration.from_pretrained(
-    "Qwen/Qwen2-VL-2B-Instruct", torch_dtype="auto",
-)
-processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
+# Load tokenizer and model for EOU detection
+tokenizer = AutoTokenizer.from_pretrained("livekit/turn-detector")
+model = AutoModelForCausalLM.from_pretrained("livekit/turn-detector")
 
-# Process text and image for inference
-def generate_response(messages: list):
-    # Preprocess conversation (text + image)
-    text_prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-    
-    # Prepare input tensors
-    images = [msg.get("image") for msg in messages if msg.get("image")]
-    text = [text_prompt]
-    
-    inputs = processor(
-        text=text, images=images, padding=True, return_tensors="pt"
-    )
+# Define function to calculate softmax
+def _softmax(logits: np.ndarray) -> np.ndarray:
+    exp_logits = np.exp(logits - np.max(logits))
+    return exp_logits / np.sum(exp_logits)
 
-    # Inference: Generate the output
-    output_ids = model.generate(**inputs, max_new_tokens=128)
-    generated_ids = [
-        output_ids[len(input_ids) :]
-        for input_ids, output_ids in zip(inputs.input_ids, output_ids)
-    ]
-    output_text = processor.batch_decode(
-        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
-    )
-    
-    return output_text[0]
+# Define the EOU probability calculation
+def get_eou_probability(chat_ctx: list) -> float:
+    """Calculate the probability of End of Utterance (EOU)"""
+    # Normalize and prepare the chat context
+    text = " ".join([msg["content"] for msg in chat_ctx])
+    inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
 
-# Gradio chat interface function
-def chat_interface(user_input, image: Image = None, history=[]):
-    # Add user input to the history
-    if image:
-        message = {
-            "role": "user",
-            "content": [{"type": "image", "image": image}, {"type": "text", "text": user_input}],
-        }
-    else:
-        message = {
-            "role": "user",
-            "content": [{"type": "text", "text": user_input}],
-        }
-    history.append(message)
+    # Run the model and get the logits
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits[0, -1, :]  # Get logits of the last token
+        probs = _softmax(logits.numpy())   # Convert logits to probabilities
 
-    # Get model response
-    response = generate_response(history)
-    
-    # Add model response to the history
-    history.append({"role": "assistant", "content": [{"type": "text", "text": response}]})
+    # Assuming <|im_end|> token corresponds to EOU, get the probability of that token
+    eou_token_id = tokenizer.encode("<|im_end|>")[-1]
+    return probs[eou_token_id]
+
+# Define the main response function for Gradio
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    messages = [{"role": "system", "content": system_message}]
     
-    return history, response
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
+
+    messages.append({"role": "user", "content": message})
+
+    # Get the response from the Qwen model (e.g., for conversation generation)
+    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
 
-# Gradio chat interface setup
-def create_gradio_interface():
-    # Chat interface with image upload and text input
-    interface = gr.Interface(
-        fn=chat_interface,
-        inputs=[
-            gr.Textbox(type="text", label="Your Message"),
-            gr.Image(type="pil", label="Upload an Image", optional=True)
-        ],
-        outputs=[
-            gr.Chatbot(label="Chatbot"),
-            gr.Textbox(label="Model's Response")
-        ],
-        title="Chat with Vision Model",
-        description="This is a multimodal model where you can chat with it using both images and text inputs. The model will respond accordingly based on your input.",
-        allow_flagging="never"
-    )
+    # After generating the response, get the EOU probability
+    eou_probability = get_eou_probability(messages)  # Get EOU prediction
+    print(f"EOU Probability: {eou_probability}")
 
-    return interface
+    # Include the EOU probability in the output
+    yield f"\nEOU Probability: {eou_probability:.2f}"
+
+# Gradio interface setup
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Textbox(value="Bạn là một trợ lý ảo", label="System message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, 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)",
+        ),
+    ],
+)
 
-# Run the Gradio app
 if __name__ == "__main__":
-    interface = create_gradio_interface()
-    interface.launch()
+    demo.launch()