my-ai-space2 / app.py
broadfield-dev's picture
Initial Space setup of broadfield-dev/my-ai-space2 via Builder
fb19138 verified
raw
history blame contribute delete
800 Bytes
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
def generate_text(prompt, length):
model_name = "EleutherAI/gpt-neo-1.3B"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer(prompt, return_tensors="pt")
output = model.generate(**inputs, max_length=length, num_beams=4, early_stopping=True)
return tokenizer.decode(output[0], skip_special_tokens=True)
demo = gr.Interface(
fn=generate_text,
inputs=["text", "int"],
outputs="text",
title="Hugging Face Text Generation",
description="Enter a prompt and a length, and this chatbot will generate text using a pre-trained GPT-Neo model."
)
if __name__ == "__main__":
demo.launch()