initial deployment

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.10-slim
+
+# HuggingFace Spaces requires non-root user with UID 1000
+RUN useradd -m -u 1000 user
+WORKDIR /home/user/app
+
+# Install dependencies first (layer caching)
+COPY --chown=user requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY --chown=user app.py .
+
+# Copy model files (large layer — copy last to keep build cache useful)
+COPY --chown=user saved_models/ ./saved_models/
+
+USER user
+
+EXPOSE 7860
+
+ENV MODEL_DIR=./saved_models/mentalbert_v3flat
+ENV TOKENIZER_DIR=./saved_models/mentalbert_v3flat/tokenizer
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,12 +1,40 @@
 ---
-title: Vibecheck Api
-emoji: 🔥
-colorFrom: blue
-colorTo: blue
+title: VibeCheck MentalBERT API
+emoji: 🧠
+colorFrom: indigo
+colorTo: purple
 sdk: docker
+app_port: 7860
 pinned: false
 license: mit
-short_description: VibeCheck NLP's API
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# VibeCheck MentalBERT API
+
+FastAPI backend serving a fine-tuned MentalBERT model for mental health text classification.
+
+## Endpoint
+
+**POST** `/classify`
+
+Request body:
+```json
+{ "text": "I have been feeling really overwhelmed lately" }
+```
+
+Response:
+```json
+{ "classification": "stress", "confidence": 0.8731 }
+```
+
+## Classes
+`anxiety` · `bipolar` · `depression` · `normal` · `personality_disorder` · `stress` · `suicidal`
+
+## Health check
+**GET** `/` → `{ "status": "ok", "model_loaded": true }`
+
+## Deployment note
+When pushing to this Space, use git-lfs for model files:
+```bash
+git lfs track "*.safetensors" "*.pt" "*.joblib"
+```

app.py

@@ -0,0 +1,120 @@
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, BertForSequenceClassification
+import joblib
+import os
+
+_BASE = os.path.dirname(os.path.abspath(__file__))
+MODEL_DIR = os.getenv(
+    "MODEL_DIR",
+    os.path.join(_BASE, "..", "saved_models", "mentalbert_v3flat")
+)
+TOKENIZER_DIR = os.getenv(
+    "TOKENIZER_DIR",
+    os.path.join(MODEL_DIR, "tokenizer")
+)
+CHECKPOINT_PATH = os.getenv(
+    "CHECKPOINT_PATH",
+    os.path.join(_BASE, "..", "saved_models", "mentalbert_v3flat_best.pt")
+)
+LABEL_ENCODER_PATH = os.path.join(MODEL_DIR, "label_encoder.joblib")
+BASE_MODEL_NAME = "mental/mental-bert-base-uncased"
+N_CLASSES = 7
+MAX_LEN = 128
+DEVICE = torch.device("cpu")
+
+LABEL_MAP: dict[str, str] = {
+    "Anxiety": "anxiety",
+    "Bipolar": "bipolar",
+    "Depression": "depression",
+    "Normal": "normal",
+    "Personality Disorder": "personality_disorder",
+    "Stress": "stress",
+    "Suicidal": "suicidal",
+}
+
+model_state: dict = {}
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
+
+    # Load architecture from HuggingFace base, then overwrite weights from .pt checkpoint
+    model = BertForSequenceClassification.from_pretrained(
+        BASE_MODEL_NAME, num_labels=N_CLASSES, ignore_mismatched_sizes=True
+    )
+    checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
+    # .pt may be a raw state_dict or wrapped under a key
+    state_dict = checkpoint.get("model_state_dict", checkpoint)
+    model.load_state_dict(state_dict)
+
+    model.to(DEVICE)
+    model.eval()
+    label_encoder = joblib.load(LABEL_ENCODER_PATH)
+    model_state.update({"tokenizer": tokenizer, "model": model, "label_encoder": label_encoder})
+    print("MentalBERT model loaded successfully.")
+    yield
+    model_state.clear()
+
+
+app = FastAPI(title="VibeCheck API", version="1.0.0", lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=False,
+    allow_methods=["POST", "GET"],
+    allow_headers=["Content-Type"],
+)
+
+
+class ClassifyRequest(BaseModel):
+    text: str
+
+
+class ClassifyResponse(BaseModel):
+    classification: str
+    confidence: float
+
+
+@app.get("/")
+def health():
+    return {"status": "ok", "model_loaded": bool(model_state)}
+
+
+@app.post("/classify", response_model=ClassifyResponse)
+def classify(req: ClassifyRequest):
+    text = req.text.strip()
+    if not text:
+        raise HTTPException(status_code=422, detail="text must not be empty")
+
+    tokenizer = model_state["tokenizer"]
+    model = model_state["model"]
+    label_encoder = model_state["label_encoder"]
+
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding="max_length",
+        max_length=MAX_LEN,
+    )
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+
+    with torch.no_grad():
+        logits = model(**inputs).logits
+
+    probs = F.softmax(logits, dim=-1)
+    confidence = float(probs.max().item())
+    pred_idx = int(torch.argmax(probs, dim=-1).item())
+    raw_label: str = label_encoder.inverse_transform([pred_idx])[0]
+
+    return ClassifyResponse(
+        classification=LABEL_MAP.get(raw_label, "normal"),
+        confidence=round(confidence, 4),
+    )

requirements.txt

@@ -0,0 +1,9 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.3.1+cpu
+transformers==4.43.4
+fastapi==0.111.0
+uvicorn[standard]==0.30.1
+scikit-learn==1.5.1
+joblib==1.4.2
+safetensors==0.4.3
+pydantic==2.8.2

saved_models/mentalbert_v3flat/config.json

@@ -0,0 +1,21 @@
+{
+  "model_name": "mental/mental-bert-base-uncased",
+  "dataset": "mental_health_v3.csv (mentaldistress excluded)",
+  "architecture": "flat 7-class",
+  "max_len": 128,
+  "n_classes": 7,
+  "classes": [
+    "Anxiety",
+    "Bipolar",
+    "Depression",
+    "Normal",
+    "Personality Disorder",
+    "Stress",
+    "Suicidal"
+  ],
+  "test_acc": 0.8197624586751561,
+  "test_f1_macro": 0.7646452685254392,
+  "test_f1_weighted": 0.8193583261268111,
+  "macro_auc": 0.9730642236990662,
+  "dep_sui_bleed": 1224
+}

saved_models/mentalbert_v3flat/label_encoder.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e0fd16abb836f540e7b52e71747fc16cc1be299c8024e51284d9cd1b99c226e6
+size 558

saved_models/mentalbert_v3flat/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1357b170b0d3335f475148307e3348a327ed5979c38447dbd049c4411e26a58
+size 437974004

saved_models/mentalbert_v3flat/tokenizer/tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "backend": "tokenizers",
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "is_local": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}

saved_models/mentalbert_v3flat_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55997f0777beac09dd99b9100e5bbcf1a03de40b40ef02e9a8efef62ecd57869
+size 438035662