import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F

# Load model and tokenizer (same as the lesson)
CHECKPOINT = "distilbert-base-uncased-finetuned-sst-2-english"
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
model = AutoModelForSequenceClassification.from_pretrained(CHECKPOINT)

def behind_the_pipeline(text):
    # Step 1: Tokenization / Preprocessing
    inputs = tokenizer(text, return_tensors="pt", truncation=True)
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    
    # Step 2: Model Inference (get logits)
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits[0]
    
    # Step 3: Post-processing (softmax for probabilities)
    probs = F.softmax(logits, dim=-1)
    id2label = model.config.id2label
    labels = [id2label[i] for i in range(len(probs))]
    
    # Build results
    tokenization_str = f"**Tokens:** {tokens}\n\n**Input IDs:** {inputs['input_ids'][0].tolist()}"
    logits_str = f"**Logits (raw model output):**\n`{logits.tolist()}`"
    probs_str = "\n".join([
        f"- **{label}**: {prob:.4f}" 
        for label, prob in zip(labels, probs.tolist())
    ])
    prediction = id2label[probs.argmax().item()]
    
    return tokenization_str, logits_str, probs_str, prediction

# Gradio UI
with gr.Blocks(title="Behind the Pipeline") as demo:
    gr.Markdown("""
    # 🕵️ Behind the Pipeline
    Replicating the [Hugging Face NLP Course](https://huggingface.co/learn/nlp-course/chapter2/2) lesson in an interactive Space.
    
    This demonstrates the **three steps** hidden inside `pipeline("sentiment-analysis")`:
    1. **Preprocessing** (Tokenizer)
    2. **Model Inference** (Logits)
    3. **Post-processing** (Softmax → Probabilities)
    """)
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Input Text", 
                value="I've been waiting for a HuggingFace course my whole life."
            )
            run_btn = gr.Button("Run Pipeline", variant="primary")
        
        with gr.Column():
            prediction = gr.Textbox(label="Final Prediction", interactive=False)
    
    with gr.Row():
        with gr.Column():
            tokenization = gr.Markdown(label="Step 1: Tokenization")
        with gr.Column():
            logits = gr.Markdown(label="Step 2: Logits")
        with gr.Column():
            probabilities = gr.Markdown(label="Step 3: Probabilities")
    
    run_btn.click(
        fn=behind_the_pipeline,
        inputs=input_text,
        outputs=[tokenization, logits, probabilities, prediction]
    )

    # Examples from the lesson
    gr.Examples(
        examples=[
            "I've been waiting for a HuggingFace course my whole life.",
            "I hate this so much!",
            "This movie was surprisingly good.",
            "The weather today is terrible."
        ],
        inputs=input_text
    )

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