app.py

import gradio as gr
from transformers import pipeline
import torch

# Load the model with custom pipeline (includes built-in rules!)
model_id = "samurai9776/thought-classifier"
classifier = pipeline(
    "text-classification", 
    model=model_id,
    trust_remote_code=True,  # Required for custom pipeline with rules
    device=0 if torch.cuda.is_available() else -1
)

def classify_thought(ai_utterance, cx_utterance):
    """Classify if the conversation is complete or incomplete"""
    
    if not ai_utterance or not cx_utterance:
        return "Please enter both utterances", 0, 0, ""
    
    # Combine utterances with [SEP] token
    text = f"{ai_utterance} [SEP] {cx_utterance}"
    
    # Get prediction from pipeline (now includes built-in rules!)
    results = classifier(text)
    
    # Extract scores
    scores = {r['label']: r['score'] for r in results}
    complete_score = scores.get('COMPLETE', 0)
    incomplete_score = scores.get('INCOMPLETE', 0)
    
    # Determine prediction
    prediction = "Complete ✓" if complete_score > incomplete_score else "Incomplete ⚠️"
    
    # Check if it's a rule-based decision (very high confidence)
    if max(complete_score, incomplete_score) >= 0.94:
        method = "Rule-based detection (linguistic pattern)"
    else:
        method = "Neural model prediction"
    
    return prediction, complete_score, incomplete_score, method

# Create Gradio interface
with gr.Blocks(
    title="Thought Completion Classifier - Professional Edition",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        font-family: 'IBM Plex Sans', sans-serif;
    }
    .gr-button {
        font-size: 16px;
    }
    """
) as demo:
    gr.Markdown("""
    # 🤖 Thought Completion Classifier
    ### Professional Edition with Linguistic Rules Engine
    
    This advanced model determines if a conversation represents a **complete** or **incomplete** thought.
    It combines linguistic rules with neural networks for superior accuracy.
    
    #### 🎯 Built-in Linguistic Rules:
    - **Incomplete patterns**: Ends with 'a', 'and', 'with', 'to', 'for', 'of', 'or', 'what', 'get'
    - **Complete patterns**: Contains 'that's all', 'nothing else', 'yes please', 'perfect'
    - **Smart fallback**: Uses neural model for complex cases
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            ai_input = gr.Textbox(
                label="AI Utterance",
                placeholder="e.g., What else can I get for you?",
                lines=2
            )
            cx_input = gr.Textbox(
                label="Customer Utterance", 
                placeholder="e.g., I need a",
                lines=2
            )
            
            with gr.Row():
                classify_btn = gr.Button("🔍 Classify", variant="primary", scale=2)
                clear_btn = gr.Button("🗑️ Clear", variant="secondary", scale=1)
            
        with gr.Column(scale=1):
            prediction = gr.Textbox(
                label="Prediction", 
                interactive=False,
                elem_classes=["prediction-box"]
            )
            with gr.Row():
                complete_score = gr.Number(
                    label="Complete Score", 
                    precision=3, 
                    interactive=False
                )
                incomplete_score = gr.Number(
                    label="Incomplete Score", 
                    precision=3, 
                    interactive=False
                )
            method = gr.Textbox(
                label="Detection Method", 
                interactive=False,
                elem_classes=["method-box"]
            )
    
    # Examples section
    gr.Markdown("### 📝 Try These Examples:")
    gr.Examples(
        examples=[
            ["What else?", "I need a", "Rule: Ends with 'a' → Incomplete"],
            ["Anything else?", "And", "Rule: Ends with 'and' → Incomplete"],
            ["Is that all?", "That's all", "Rule: Contains 'that's all' → Complete"],
            ["Your order?", "Perfect", "Rule: Contains 'perfect' → Complete"],
            ["Can I help you?", "I want to", "Rule: Ends with 'to' → Incomplete"],
            ["What would you like?", "I'll have the burger please", "Model: Complex sentence"],
            ["Ready to checkout?", "Actually let me add", "Rule: Contains 'let me add' → Incomplete"],
            ["Anything else today?", "No thanks", "Rule: Contains 'no thanks' → Complete"],
        ],
        inputs=[ai_input, cx_input],
        outputs=[prediction, complete_score, incomplete_score, method],
        fn=classify_thought,
        cache_examples=True,
        label="Click any example to test it"
    )
    
    # Event handlers
    classify_btn.click(
        fn=classify_thought,
        inputs=[ai_input, cx_input],
        outputs=[prediction, complete_score, incomplete_score, method]
    )
    
    clear_btn.click(
        fn=lambda: ("", "", "", 0, 0, ""),
        inputs=[],
        outputs=[ai_input, cx_input, prediction, complete_score, incomplete_score, method]
    )
    
    # API usage section
    with gr.Accordion("🔧 API Usage", open=False):
        gr.Markdown("""
        ### Use this model in your code:
        ```python
        from transformers import pipeline
        
        # Load the classifier with built-in rules
        classifier = pipeline(
            "text-classification", 
            model="samurai9776/thought-classifier",
            trust_remote_code=True  # Required for rules engine
        )
        
        # Example usage
        result = classifier("What else? [SEP] I need a")
        print(result)  # [{'label': 'INCOMPLETE', 'score': 0.95}]
        ```
        
        ### Model Information:
        - **Base Model**: DistilBERT (66M parameters)
        - **Enhancement**: Custom pipeline with linguistic rules
        - **Training Data**: 900+ labeled conversations
        - **Accuracy**: 90%+ with rules engine
        """)
    
    # Footer
    gr.Markdown("""
    ---
    <div style='text-align: center; color: #666;'>
        <p>Model by <a href='https://huggingface.co/samurai9776' target='_blank'>samurai9776</a> | 
        <a href='https://huggingface.co/samurai9776/thought-classifier' target='_blank'>View Model</a> | 
        Built with 🤗 Transformers</p>
    </div>
    """)

if __name__ == "__main__":
    demo.launch()