import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("your-hf-username/your-model-name")
model = AutoModelForCausalLM.from_pretrained("your-hf-username/your-model-name")

# Define function to generate responses
def generate_response(persona, prompt, max_length=100):
    full_prompt = f"You are {persona}. {prompt}"
    inputs = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True).to("cuda")
    outputs = model.generate(inputs.input_ids, max_length=max_length, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Gradio interface
examples = [
    ["a grizzled pirate afraid of sharks", "What do you think of the open sea?"],
    ["a curious scientist exploring Mars", "What do you think about the red planet?"],
    ["a medieval knight on a quest for treasure", "How do you feel about dragons?"]
]

interface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(label="Describe your persona", placeholder="e.g., a grizzled pirate afraid of sharks"),
        gr.Textbox(label="Enter your prompt", placeholder="e.g., What do you think of the open sea?"),
        gr.Slider(10, 200, value=100, step=10, label="Response Length")
    ],
    outputs="text",
    examples=examples,
    title="Custom Persona Chatbot",
    description="Create a persona and interact with the chatbot. Describe the persona and ask any question!"
)

interface.launch()