File size: 2,732 Bytes
974290d db03fc1 974290d db03fc1 974290d cd578be 974290d d073467 974290d 519fc13 974290d db03fc1 974290d db03fc1 974290d | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | import os
import re
from contextlib import asynccontextmanager
from typing import Literal
import torch
from fastapi import FastAPI
from pydantic import BaseModel, Field
from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer
MODEL_NAME = os.getenv("MODEL_NAME", "distilbert-base-multilingual-cased")
MODEL_WEIGHTS_PATH = os.getenv("MODEL_WEIGHTS_PATH", "depression_model.pt")
TOKENIZER_PATH = os.getenv("TOKENIZER_PATH", "tokenizer")
_device: torch.device | None = None
_tokenizer = None
_model = None
def clean_text(text: str) -> str:
text = str(text).lower()
text = re.sub(r"http\S+", "", text) # remove urls
text = re.sub(r"@\w+", "", text) # remove mentions
text = re.sub(r"[^\w\s]", "", text) # remove emojis + punctuation
text = re.sub(r"\s+", " ", text) # remove extra spaces
return text.strip()
def _load_model() -> None:
global _device, _tokenizer, _model
_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer_source = TOKENIZER_PATH if os.path.isdir(TOKENIZER_PATH) else MODEL_NAME
_tokenizer = AutoTokenizer.from_pretrained(tokenizer_source)
# Avoid downloading large base-model weights. We only need the config to
# construct the architecture, then we load your local trained weights.
config = AutoConfig.from_pretrained(MODEL_NAME, num_labels=2)
_model = AutoModelForSequenceClassification.from_config(config)
state_dict = torch.load(MODEL_WEIGHTS_PATH, map_location="cpu")
_model.load_state_dict(state_dict)
_model.to(_device)
_model.eval()
@asynccontextmanager
async def lifespan(_: FastAPI):
_load_model()
yield
app = FastAPI(title="Depression Text Classifier", version="1.0.0", lifespan=lifespan)
class PredictRequest(BaseModel):
text: str = Field(..., min_length=1, description="Input text to classify")
class PredictResponse(BaseModel):
label: Literal["Depressed", "Non-Depressed"]
class_id: int
@app.get("/health")
def health() -> dict:
return {"status": "ok"}
@app.post("/predict", response_model=PredictResponse)
def predict(req: PredictRequest) -> PredictResponse:
if _model is None or _tokenizer is None or _device is None:
_load_model()
text = clean_text(req.text)
inputs = _tokenizer(
text,
return_tensors="pt",
truncation=True,
padding=True,
max_length=128,
)
inputs = {k: v.to(_device) for k, v in inputs.items()}
with torch.no_grad():
outputs = _model(**inputs)
pred = int(torch.argmax(outputs.logits, dim=1).item())
label = "Depressed" if pred == 1 else "Non-Depressed"
return PredictResponse(label=label, class_id=pred)
|