File size: 2,090 Bytes
ad9d86e 77ba456 ad9d86e 77ba456 ad9d86e 77ba456 ad9d86e 77ba456 ad9d86e 77ba456 ad9d86e 77ba456 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from huggingface_hub import hf_hub_download
import torch
import json
MODEL_DIR = "Fatima1412/ingredients-distilbert-classifier"
tokenizer = None
model = None
id2label = None
app = FastAPI(title="Ingredients Classifier API")
@app.on_event("startup")
def load_model():
global tokenizer, model, id2label
try:
# Load tokenizer + model from Hugging Face
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
model.eval()
# Download labels_mapping.json from Hugging Face
mapping_path = hf_hub_download(
repo_id=MODEL_DIR,
filename="labels_mapping.json"
)
with open(mapping_path, "r") as f:
id2label = {int(k): v for k, v in json.load(f).items()}
print("Model and label mapping loaded successfully.")
except Exception as e:
print("⚠ WARNING: Model not loaded. Check Hugging Face repo.")
print(e)
class IngredientsRequest(BaseModel):
ingredients: str
class PredictionResponse(BaseModel):
label_id: int
label_name: str
probabilities: List[float]
def predict_ingredients(text):
inputs = tokenizer(text, return_tensors="pt")
# DistilBERT does NOT accept token_type_ids
if "token_type_ids" in inputs:
del inputs["token_type_ids"]
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probs = torch.softmax(logits, dim=-1).tolist()[0]
label_id = logits.argmax(dim=-1).item()
label_name = id2label[label_id]
return label_id, label_name, probs
@app.post("/predict", response_model=PredictionResponse)
def predict(req: IngredientsRequest):
label_id, label_name, probs = predict_ingredients(req.ingredients)
return PredictionResponse(
label_id=label_id,
label_name=label_name,
probabilities=probs
)
|