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()