import time
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr

# Load model and tokenizer once at startup
model_name = "Qwen/Qwen2.5-0.5B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
print(f"Model loaded on: {model.device}")

# Define the generation function
def generate_text(prompt, max_new_tokens, num_beams):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    input_length = inputs["input_ids"].shape[-1]  
    # Greedy search
    start_time = time.time()
    outputs_greedy = model.generate(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        num_beams=1,
        do_sample=False,
    )
    greedy_time = time.time() - start_time
    # Remove the prompt tokens from the output
    generated_tokens_greedy = outputs_greedy[0][input_length:]
    generated_text_greedy = tokenizer.decode(generated_tokens_greedy, skip_special_tokens=True)
    
    # Beam search
    start_time = time.time()
    outputs_beam = model.generate(
        **inputs,
        num_beams=int(num_beams),
        num_return_sequences=1,
        max_new_tokens=int(max_new_tokens),
        do_sample=False,
    )
    beam_time = time.time() - start_time
    # Remove the prompt tokens as above
    generated_tokens_beam = outputs_beam[0][input_length:]
    generated_text_beam = tokenizer.decode(generated_tokens_beam, skip_special_tokens=True)
    
    # Prepare outputs for better display formatting
    greedy_details = (
        f"**Strategy:** Picks the most probable token at each step (deterministic).\n\n"
        f"**Time:** {greedy_time:.2f} seconds"
    )
    
    beam_details = (
        f"**Strategy:** Explores {num_beams} beams concurrently and returns the top candidate.\n\n"
        f"**Time:** {beam_time:.2f} seconds"
    )
    
    return greedy_details, generated_text_greedy, beam_details, generated_text_beam

with gr.Blocks() as demo:
    # Informational header to help users understand the demo
    gr.Markdown(
        "# Beam Search Demo\n\n"
        "This demo shows how two different text generation strategies work using the Qwen2.5-0.5B model. "
        "The left side uses **greedy search**, which picks the most probable token at every generation step (deterministic), "
        "while the right side uses **beam search**, which explores multiple beams concurrently to choose the most likely "
        "sequence of tokens.\n\n"
        "**Important:** This model works best with prompts that need completion rather than question-answering. For example, "
        "instead of 'What is the capital of France?', use prompts like 'The capital of France is' or 'Here is a story about:'\n\n"
        "Use the controls below to enter your prompt, adjust the maximum number of newly generated tokens, and set the "
        "number of beams for beam search. The results for both strategies are displayed side by side for easy comparison.\n\n"
        "Repo: [Beam Search Demo](https://github.com/cavit99/beam-search-demo)"
    )

    # Input components in a single column at the top
    with gr.Column():
        gr.Markdown("## Input")
        prompt_input = gr.Textbox(label="Prompt", value="Here is a funny love letter for you:")
        max_tokens_input = gr.Slider(minimum=1, maximum=100, step=1, label="Max new tokens", value=50)
        num_beams_input = gr.Slider(minimum=1, maximum=20, step=1, label="Number of beams", value=10)
        generate_btn = gr.Button("Generate")
    
    with gr.Row():
        with gr.Column():
            greedy_details_output = gr.Markdown(label="Greedy Search Details")
            greedy_textbox_output = gr.Textbox(label="Greedy Search Generated Text", lines=10)
        with gr.Column():
            beam_details_output = gr.Markdown(label="Beam Search Details")
            beam_textbox_output = gr.Textbox(label="Beam Search Generated Text", lines=10)
    
    # Connect the button click event to the generation function
    generate_btn.click(
        generate_text,
        inputs=[prompt_input, max_tokens_input, num_beams_input],
        outputs=[greedy_details_output, greedy_textbox_output, beam_details_output, beam_textbox_output]
    )

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