mohameddshaheer/field-classification-api

1import json
2import torch
3from fastapi import FastAPI, HTTPException
4from fastapi.middleware.cors import CORSMiddleware
5from pydantic import BaseModel
6from transformers import AutoTokenizer, AutoModelForSequenceClassification
7
8MODEL_PATH = "./final_model"   # ← points to folder inside Docker
9TOP_K = 3
10MAX_LENGTH = 512
11
12print("Loading model...")
13tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
14model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
15model.eval()
16
17with open(f"{MODEL_PATH}/label_mapping.json", "r", encoding="utf-8") as f:
18    mapping = json.load(f)
19id2label = mapping["id2label"]
20print("Model loaded. API ready!")
21
22app = FastAPI(title="Field Classification API")
23
24app.add_middleware(
25    CORSMiddleware,
26    allow_origins=["*"],
27    allow_methods=["*"],
28    allow_headers=["*"],
29)
30
31class PredictRequest(BaseModel):
32    text: str
33
34class Prediction(BaseModel):
35    rank: int
36    label: str
37    confidence_percent: float
38
39class PredictResponse(BaseModel):
40    predictions: list[Prediction]
41    confidence_sum_percent: float
42    input_text_preview: str
43
44@app.get("/")
45def root():
46    return {"status": "ok", "message": "Field Classification API is running!"}
47
48@app.post("/predict", response_model=PredictResponse)
49def predict(request: PredictRequest):
50    text = request.text.strip()
51    if not text:
52        raise HTTPException(status_code=400, detail="text must not be empty.")
53
54    inputs = tokenizer(
55        text,
56        return_tensors="pt",
57        truncation=True,
58        padding=True,
59        max_length=MAX_LENGTH,
60    )
61
62    with torch.no_grad():
63        outputs = model(**inputs)
64
65    logits = outputs.logits[0]
66    top_values, top_indices = torch.topk(logits, TOP_K)
67    percentages = torch.softmax(top_values, dim=0) * 100
68
69    predictions = [
70        Prediction(
71            rank=i + 1,
72            label=id2label[str(idx.item())],
73            confidence_percent=round(pct.item(), 2),
74        )
75        for i, (idx, pct) in enumerate(zip(top_indices, percentages))
76    ]
77
78    return PredictResponse(
79        predictions=predictions,
80        confidence_sum_percent=round(percentages.sum().item(), 2),
81        input_text_preview=text[:120],
82    )