Shakespeare-GPT / app.py
mstftmk's picture
Add all files
91d390a
raw
history blame
1.05 kB
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load the model and tokenizer
model_name = "mstftmk/shakespeare-gpt2" # Replace with your model's repo
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Define the generation function
def generate_text(input_text, max_length, temperature):
inputs = tokenizer.encode(input_text, return_tensors="pt")
outputs = model.generate(inputs, max_length=max_length, temperature=temperature)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Create the Gradio interface
interface = gr.Interface(
fn=generate_text,
inputs=[
gr.Textbox(lines=2, placeholder="Enter your prompt..."),
gr.Slider(50, 300, value=100, label="Max Length"),
gr.Slider(0.1, 1.0, value=0.7, label="Temperature"),
],
outputs=gr.Textbox(),
title="Shakespeare GPT-2",
description="Generate text inspired by Shakespeare.",
)
# Launch the Gradio app
interface.launch()