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

# Load model and tokenizer
model_path = "ibm-granite/granite-3.0-1b-a400m-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
model.eval()

def generate_response(prompt, max_new_tokens, temperature, top_p, repetition_penalty):
    chat = [
        {"role": "user", "content": prompt},
    ]
    chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
    
    input_tokens = tokenizer(chat, return_tensors="pt").to(model.device)
    
    output = model.generate(
        **input_tokens,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True
    )
    
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response.split("Human:", 1)[0].strip()

with gr.Blocks() as demo:
    gr.Markdown("# 🙋🏻‍♂️Welcome to 🌟Tonic's🪨Granite-3.0-1B-A400M-Instruct Demo")
    gr.Markdown("Enter a prompt and adjust generation parameters to interact with the 🪨Granite-3.0-1B-A400M-Instruct model.")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=5)
            generate_button = gr.Button("Generate Response")
            max_new_tokens = gr.Slider(minimum=1, maximum=500, value=100, step=1, label="Max New Tokens")
            temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top P")
            repetition_penalty = gr.Slider(minimum=1.0, maximum=2.0, value=1.1, step=0.1, label="Repetition Penalty")
    
        with gr.Column() :
            output = gr.Textbox(label="🪨Granite3-1B", lines=10)
    
    generate_button.click(
        generate_response,
        inputs=[prompt, max_new_tokens, temperature, top_p, repetition_penalty],
        outputs=output
    )
    
    gr.Markdown("## Examples")
    examples = gr.Examples(
        examples=[
            ["Tell me about the history of artificial intelligence.", 200, 0.7, 0.9, 1.1],
            ["Write a short story about a robot learning to paint.", 300, 0.8, 0.95, 1.2],
            ["Explain the concept of quantum computing to a 10-year-old.", 150, 0.6, 0.85, 1.0],
        ],
        inputs=[prompt, max_new_tokens, temperature, top_p, repetition_penalty],
    )

demo.launch()