amarvm/BATMAN

1from fastapi import FastAPI2from pydantic import BaseModel3from transformers import pipeline4 5app = FastAPI()6generator = pipeline("text-generation", model="sshleifer/tiny-gpt2")7 8class InputText(BaseModel):9    prompt: str10 11@app.get("/")12def home():13    return {"message": "Welcome to my Hugging Face Docker App!"}14 15@app.post("/predict")16def predict(data: InputText):17    result = generator(data.prompt, max_length=40)18    return {"result": result[0]["generated_text"]}19