import gradio as gr
from transformers import pipeline, AutoTokenizer

# 1. Load the model/pipeline (runs once at startup)
classifier = pipeline("sentiment-analysis")

# 2. Define the function Gradio will call for pipeline
def analyze(text):
    # The pipeline can take a single string or a list.
    # Here we wrap a single user input in a list to match your example.
    results = classifier([text])
    
    # Pretty-print the result
    r = results[0]
    return f"Label: {r['label']}\nConfidence: {r['score']:.4f}"

# Build the UI
# demo = gr.Interface(
#     fn=analyze,                                     # your pipeline wrapper
#     inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
#     outputs=gr.Textbox(label="Result"),
#     title="Sentiment Analysis",
#     description="Running `transformers.pipeline('sentiment-analysis')` behind the scenes.",
#     examples=[
#         ["I've been waiting for a HuggingFace course my whole life."],
#         ["I hate this so much!"]
#     ]
# )

# Load tokenizer once (outside the UI)
checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

def tokenize_input(raw_inputs):
    """Tokenize text and return displayable results."""
    if not raw_inputs or not raw_inputs.strip():
        return {}, "Please enter some text."
    
    inputs = tokenizer(raw_inputs, padding=True, truncation=True, return_tensors="pt")

    # Convert tensors to plain Python types so Gradio can display them
    result = {
        "input_ids": inputs["input_ids"].tolist(),
        "attention_mask": inputs["attention_mask"].tolist()
    }

    # Also get human-readable token strings
    tokens = tokenizer.convert_ids_to_tokens(result["input_ids"][0])

    return result, str(tokens)
    
# ── Build the tabbed UI ─────────────────────────────────
with gr.Blocks(title="Lesson Steps Comparison") as demo:
    gr.Markdown("# 🤗 Transformers Lesson: Side-by-Side Comparison")

        # ── Tab 1 ──
    with gr.Tab("Step 1: Sentiment Analysis"):
        gr.Markdown("**`pipeline('sentiment-analysis')`**")
        with gr.Row():
            with gr.Column():
                inp1 = gr.Textbox(
                    lines=3,
                    label="Input",
                    placeholder="I've been waiting for a HuggingFace course..."
                )
                btn1 = gr.Button("Analyze", variant="primary")
            out1 = gr.Textbox(label="Sentiment Result")
            btn1.click(analyze, inputs=inp1, outputs=out1)

        gr.Examples(
            examples=[
                "I've been waiting for a HuggingFace course my whole life.",
                "I hate this so much!"
            ],
            inputs=inp1
        )
        # ── Tab 2 ──
    with gr.Tab("Tokenize"):
        gr.Markdown("### DistilBERT Tokenizer")
        with gr.Row():
            with gr.Column():
                text_input = gr.Textbox(lines=3, label="Input", placeholder="I've been waiting for a HuggingFace course...")
                tokenize_btn = gr.Button("Tokenize", variant="primary")
            json_output = gr.JSON(label="Tokenized Tensor Data")
            tokens_output = gr.Textbox(label="Human-Readable Tokens", 
                lines=6)
              # Wire button → function → outputs
        tokenize_btn.click(
            fn=tokenize_input,
            inputs=text_input,
            outputs=[json_output, tokens_output]
        )

    # ── Tab 3 ──
    # with gr.Tab("Step 3: [Another Comparison]"):
    #     gr.Markdown("*Another tab for easy comparison.*")
    #     inp3 = gr.Textbox(lines=3, label="Input")
    #     btn3 = gr.Button("Run Step 3")
    #     out3 = gr.Textbox(label="Output")
    #     btn3.click(run_step_3, inputs=inp3, outputs=out3)
# 4. Launch
if __name__ == "__main__":
    demo.launch()