File size: 476 Bytes
92ddabd 70f21b9 92ddabd | 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()
generator = pipeline("text-generation", model="sshleifer/tiny-gpt2")
class InputText(BaseModel):
prompt: str
@app.get("/")
def home():
return {"message": "Welcome to my Hugging Face Docker App!rahul"}
@app.post("/predict")
def predict(data: InputText):
result = generator(data.prompt, max_length=40)
return {"result": result[0]["generated_text"]}
|