tpu / app.py
mohab94's picture
Create app.py
80edb37 verified
raw
history blame contribute delete
611 Bytes
import gradio as gr
import spaces
from transformers import pipeline
# Load a small model that runs fast
generator = pipeline("text-generation", model="gpt2")
@spaces.GPU
def generate(text):
result = generator(text, max_length=50, num_return_sequences=1)
return result[0]['generated_text']
# Create the UI
with gr.Blocks() as demo:
gr.Markdown("# ๐Ÿš€ My Free H200 Space")
input_text = gr.Textbox(label="Enter a prompt")
output_text = gr.Textbox(label="Generated Text")
btn = gr.Button("Generate")
btn.click(fn=generate, inputs=input_text, outputs=output_text)
demo.launch()