import os
from fastapi import FastAPI
from transformers import pipeline

# Configurar el directorio de caché de Hugging Face
os.environ['TRANSFORMERS_CACHE'] = '/tmp/cache'

# Crear el directorio de caché si no existe
os.makedirs('/tmp/cache', exist_ok=True)

app = FastAPI()

# Cargar el modelo de GPT-Neo
generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B')

@app.post("/generate")
async def generate(prompt: str):
    response = generator(prompt, max_length=100)
    return {"text": response[0]["generated_text"]}