|
|
from fastapi import FastAPI, HTTPException |
|
|
import uvicorn |
|
|
from typing import Union, List |
|
|
from predict_dept_model import DepartmentPredictor |
|
|
from contextlib import asynccontextmanager |
|
|
from response_schema import ClassificationOutput, TextInput |
|
|
from huggingface_hub import HfApi |
|
|
import os |
|
|
|
|
|
|
|
|
model_repo = os.getenv("MODEL_REPO") |
|
|
|
|
|
|
|
|
api = HfApi() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@asynccontextmanager |
|
|
async def lifespan(app: FastAPI): |
|
|
|
|
|
global predictor |
|
|
predictor = DepartmentPredictor(model_repo= model_repo) |
|
|
yield |
|
|
|
|
|
|
|
|
app = FastAPI( |
|
|
title="Sambodhan Department Classifier API", |
|
|
description="AI model that classifies citizen grievances into municipal departments with confidence scores.", |
|
|
version="1.0.0", |
|
|
lifespan=lifespan |
|
|
) |
|
|
|
|
|
|
|
|
@app.post("/predict", response_model=Union[ClassificationOutput, List[ClassificationOutput]]) |
|
|
def predict_department(input_data: TextInput): |
|
|
try: |
|
|
|
|
|
prediction = predictor.predict(input_data.text) |
|
|
|
|
|
|
|
|
return prediction |
|
|
|
|
|
except Exception as e: |
|
|
|
|
|
raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}") |
|
|
|
|
|
|
|
|
|
|
|
@app.get("/") |
|
|
def root(): |
|
|
|
|
|
latest_tag = api.list_repo_refs(repo_id=model_repo, repo_type="model").tags[0].name |
|
|
|
|
|
return { |
|
|
"message": "Sambodhan Department Classification API is running.", |
|
|
"status": "Active" if predictor else "Inactive", |
|
|
"model_version": latest_tag |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|