import gradio as gr
from transformers import pipeline

# Load the clean v2 model
classifier = pipeline(
    "text-classification", 
    model="samurai9776/thought-classifier-menu-aware"
)

def classify_thought(ai_utterance, cx_utterance):
    if not ai_utterance or not cx_utterance:
        return "Please enter both utterances", 0, 0
    
    text = f"{ai_utterance} [SEP] {cx_utterance}"
    results = classifier(text)
    
    scores = {r['label']: r['score'] for r in results}
    complete_score = scores.get('COMPLETE', 0)
    incomplete_score = scores.get('INCOMPLETE', 0)
    
    prediction = "Complete ✓" if complete_score > incomplete_score else "Incomplete ⚠️"
    return prediction, complete_score, incomplete_score

with gr.Blocks(title="Thought Classifier v2") as demo:
    gr.Markdown("# 🤖 Thought Completion Classifier v2")
    
    with gr.Row():
        with gr.Column():
            ai_input = gr.Textbox(label="AI Utterance", placeholder="What sandwich?")
            cx_input = gr.Textbox(label="Customer Utterance", placeholder="Chicken")
            classify_btn = gr.Button("Classify", variant="primary")
            
        with gr.Column():
            prediction = gr.Textbox(label="Prediction")
            complete_score = gr.Number(label="Complete Score")
            incomplete_score = gr.Number(label="Incomplete Score")
    
    classify_btn.click(
        classify_thought,
        inputs=[ai_input, cx_input],
        outputs=[prediction, complete_score, incomplete_score]
    )

demo.launch()