File size: 452 Bytes
799b050 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
app = FastAPI()
pipe = pipeline("text2text-generation", model="google/flan-t5-base")
class Req(BaseModel):
inputs: str
@app.get("/health")
def health():
return {"status": "ok"}
@app.post("/generate")
def generate(req: Req):
out = pipe(req.inputs, max_new_tokens=64)
text = out[0].get("generated_text", "").strip()
return {"text": text} |