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"]}