import gradio as gr
from transformers import pipeline
import json

# Load your model
print("Loading model...")
classifier = pipeline(
    "text-classification",
    model="LokeshDevCreates/tone-baseline-v3",
    top_k=None  # Return all labels with scores
)
print("Model loaded successfully!")

def classify_tone(text):
    """Classify tone of input text"""
    try:
        results = classifier(text)[0]
        # Sort by score descending
        results = sorted(results, key=lambda x: x['score'], reverse=True)
        
        # Return as dict for easy JSON parsing
        return {
            "detected_tone": results[0]['label'],
            "confidence": round(results[0]['score'], 4),
            "all_probs": {r['label']: round(r['score'], 4) for r in results}
        }
    except Exception as e:
        return {"error": str(e)}

# Create Gradio interface
demo = gr.Interface(
    fn=classify_tone,
    inputs=gr.Textbox(
        label="Text to Analyze",
        placeholder="Enter text here...",
        lines=3
    ),
    outputs=gr.JSON(label="Tone Analysis"),
    title="Tone Detection API",
    description="Detect the tone of text using tone-baseline-v3 model",
    examples=[
        ["This is absolutely terrible and I hate it!"],
        ["Thank you so much for your help!"],
        ["The meeting is scheduled for 3pm tomorrow."],
    ],
    api_name="predict"  # Explicitly name the API endpoint
)

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