checkin_model / app.py
Dc-4nderson's picture
Update app.py
913861e verified
raw
history blame contribute delete
838 Bytes
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
# FastAPI app
app = FastAPI(title="Check-In Expansion API")
@app.get("/")
def read_root():
return {"status": "ok"}
# Load model
generator = pipeline(
"summarization", # use summarization since text2text-generation is deprecated in HF metadata
model="Dc-4nderson/checkin-model",
tokenizer="Dc-4nderson/checkin-model"
)
# Request/Response schema
class CheckInRequest(BaseModel):
task: str
class CheckInResponse(BaseModel):
checkin: str
@app.post("/predict", response_model=CheckInResponse)
def predict(request: CheckInRequest):
"""Expand a short input into a detailed check-in."""
result = generator(request.task, max_length=200, num_beams=4, do_sample=False)
return {"checkin": result[0]["summary_text"]}