sbicy's picture
Update app.py
1530d5e verified
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()