import os
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from huggingface_hub import login

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-7B")
model = AutoModelForCausalLM.from_pretrained(
    "Zyphra/Zamba2-7B",
    device_map="cuda",  # Automatically handles device placement
    torch_dtype=torch.bfloat16
)

# Define the function to generate responses
def generate_response(input_text, max_new_tokens, temperature, top_k, top_p, repetition_penalty, num_beams, length_penalty):
    # Tokenize and move input to model's device
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(model.device)
    
    # Generate response using specified parameters
    outputs = model.generate(
        input_ids=input_ids,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        num_beams=num_beams,
        length_penalty=length_penalty,
        num_return_sequences=1
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Create Gradio interface with adjustable parameters
demo = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(lines=1, placeholder="Enter a text to prepend...", label="Input Text"),
        gr.Slider(50, 1000, step=50, value=500, label="Max New Tokens"),
        gr.Slider(0.1, 1.5, step=0.1, value=0.7, label="Temperature"),
        gr.Slider(1, 100, step=1, value=50, label="Top K"),
        gr.Slider(0.1, 1.0, step=0.1, value=0.9, label="Top P"),
        gr.Slider(1.0, 2.0, step=0.1, value=1.2, label="Repetition Penalty"),
        gr.Slider(1, 10, step=1, value=5, label="Number of Beams"),
        gr.Slider(0.0, 2.0, step=0.1, value=1.0, label="Length Penalty")
    ],
    outputs=gr.Textbox(label="Generated Response"),
    title="Zamba2-7B Model",
    description="Ask Zamba2 7B a question with customizable parameters."
)

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