import gradio as gr
from transformers import pipeline

# ==========================================
# 1. LOAD MODEL A: Toxicity Detector (Transformer)
# ==========================================
# Uses a lightweight, high-performance BERT model for hate speech/toxicity
toxic_pipeline = pipeline("text-classification", model="unitary/toxic-bert")

def predict_toxicity(text):
    result = toxic_pipeline(text)[0]
    label = result['label']
    score = result['score']
    
    # unitary/toxic-bert outputs specific toxic labels or 'toxic' score
    if label == "toxic" and score > 0.5:
        return f"🚨 TOXIC CONTENT DETECTED! (Confidence: {score:.2f})"
    else:
        # If score is high on non-toxic aspects or label is safe
        return f"✅ Clean / Safe Content"

# ==========================================
# 2. LOAD MODEL B: Sarcasm Detector (Your Trained Model)
# ==========================================
# Loads your fine-tuned model files sitting in your current directory "."
sarcasm_pipeline = pipeline("text-classification", model=".", tokenizer=".")

def predict_sarcasm(text):
    result = sarcasm_pipeline(text)[0]
    label = result['label']
    score = result['score']
    
    if label == "LABEL_1":
        return f"😏 Sarcastic / Passive-Aggressive (Confidence: {score:.2f})"
    else:
        return f"😇 Genuine / Normal Text (Confidence: {score:.2f})"

# ==========================================
# 3. COMBINED ANALYSIS FUNCTION
# ==========================================
def analyze_text(text):
    # Run through both modern transformer models
    toxicity_result = predict_toxicity(text)
    sarcasm_result = predict_sarcasm(text)
    
    return toxicity_result, sarcasm_result

# ==========================================
# 4. GRADIO UI SETUP
# ==========================================
demo = gr.Interface(
    fn=analyze_text,
    inputs=gr.Textbox(lines=3, placeholder="Type something to analyze both Toxicity and Sarcasm..."),
    outputs=[
        gr.Textbox(label="Model 1: Toxicity Check (Hate-BERT Transformer)"),
        gr.Textbox(label="Model 2: Sarcasm Check (Your Fine-Tuned Transformer)")
    ],
    title="🛡️ Super-Duper Multi-Task Text Analyzer",
    description="This advanced interface runs your input text through two distinct Transformer pipelines simultaneously."
)

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