TextPlayground / app.py
Tralalabs's picture
Update app.py
9cb422a verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
model_id = "EleutherAI/gpt-neo-125M"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
def generate_text(prompt, temperature, max_new_tokens):
inputs = tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=int(max_new_tokens),
do_sample=True,
temperature=temperature,
top_p=0.95,
repetition_penalty=1.1,
pad_token_id=tokenizer.eos_token_id
)
full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
generated = full_text[len(prompt):]
# 🔥 combine prompt + green continuation
return f"""
<div style="font-family:monospace; white-space:pre-wrap;">
{prompt}<span style="color:#00ff88;">{generated}</span>
</div>
"""
with gr.Blocks(css="""
textarea {font-family: monospace;}
""") as demo:
gr.Markdown("# 🧠 TextPlayground")
with gr.Row():
with gr.Column(scale=3):
prompt = gr.Textbox(
label="Prompt",
placeholder="Type something like: Three reasons to start a succulent garden",
lines=10
)
output = gr.HTML(label="Output")
btn = gr.Button("Submit ⚡")
with gr.Column(scale=1):
temperature = gr.Slider(0.1, 1.5, value=0.7, label="Temperature")
max_tokens = gr.Slider(10, 256, value=120, label="Max length")
btn.click(
fn=generate_text,
inputs=[prompt, temperature, max_tokens],
outputs=output
)
demo.launch()