import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the model once
model_name = "HuggingFaceTB/SmolLM-1.7B"
model = AutoModelForCausalLM.from_pretrained(model_name)

# Define a list of five different tokenizers to use
tokenizer_names = [
    "HuggingFaceTB/SmolLM-1.7B",  # Model's default tokenizer
    "gpt2",                       # GPT-2 tokenizer
    "distilbert-base-uncased",    # DistilBERT tokenizer
    "bert-base-uncased",          # BERT tokenizer
    "roberta-base"                # RoBERTa tokenizer
]

# Load all the tokenizers
tokenizers = {name: AutoTokenizer.from_pretrained(name) for name in tokenizer_names}

# Function to generate responses using different tokenizers
def generate_responses(prompt):
    responses = {}
    for name, tokenizer in tokenizers.items():
        inputs = tokenizer(prompt, return_tensors="pt")
        outputs = model.generate(**inputs)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        responses[name] = response
    return responses

# Gradio interface setup
interface = gr.Interface(
    fn=generate_responses,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your prompt here..."),
    outputs=gr.outputs.JSON(),
    title="Tokenizer Comparison",
    description="Compare model outputs with different tokenizers"
)

# Launch the Gradio interface
interface.launch()