File size: 451 Bytes
17dd8ac
0f62da2
17dd8ac
 
0f62da2
17dd8ac
 
 
0f62da2
17dd8ac
 
 
 
 
 
0f62da2
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from fastapi import FastAPI, Request
from transformers import pipeline

app = FastAPI()
generator = pipeline("text-generation", model="gpt2")

@app.get("/")
def root():
    return {"status": "Space is running 🚀"}

@app.post("/ask")
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']}