|
import gradio as gr |
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("your-hf-username/your-model-name") |
|
model = AutoModelForCausalLM.from_pretrained("your-hf-username/your-model-name") |
|
|
|
|
|
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 |
|
|
|
|
|
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() |