| from fastapi import FastAPI, Request | |
| from transformers import pipeline | |
| app = FastAPI() | |
| generator = pipeline("text-generation", model="gpt2") | |
| def root(): | |
| return {"status": "Space is running 🚀"} | |
| async def ask(req: Request): | |
| data = await req.json() | |
| prompt = data.get("prompt", "") | |
| result = generator(prompt, max_length=100, do_sample=True) | |
| return {"generated_text": result[0]['generated_text']} | |