""" FastAPI backend for the Hierarchical Longformer. Deployed on Hugging Face Spaces (Docker SDK). Loads model from HF Hub and exposes a /predict endpoint. """ import os import re import warnings from contextlib import asynccontextmanager import torch import torch.nn as nn import torch.nn.functional as F from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from transformers import LongformerTokenizerFast, LongformerModel, LongformerConfig from huggingface_hub import hf_hub_download warnings.filterwarnings("ignore") # ── Config ─────────────────────────────────────────────────────────────────── # Change these to match your HF Hub repo HF_MODEL_REPO = "sweebano/hierarchical-longformer" MODEL_FILENAME = "hierarchical_model.pt" TOKENIZER_DIR = "hierarchical_model_tokenizer" HF_TOKEN = os.environ.get("HF_TOKEN", None) # Set as Space secret for private repos LABEL_NAMES = ["Exc", "Imp", "Imp Act", "Opp", "Opp Act", "Men", "None"] _TAG_RE = re.compile( r"\[/?(?:Exc|Imp|Imp Act|Opp|Opp Act|Men)\]", re.IGNORECASE, ) def strip_label_tags(text): return _TAG_RE.sub("", text).strip() # ── Model Definition (must match train_var.py exactly) ─────────────────────── class FocalLoss(nn.Module): def __init__(self, alpha=None, gamma=2.0, reduction="mean"): super().__init__() if alpha is not None: self.register_buffer("alpha", alpha.clone().detach()) else: self.alpha = None self.gamma = gamma self.reduction = reduction def forward(self, logits, targets): weight = self.alpha.to(logits.device) if self.alpha is not None else None ce_loss = F.cross_entropy(logits, targets, weight=weight, reduction="none") pt = torch.exp(-ce_loss) focal_loss = ((1 - pt) ** self.gamma) * ce_loss if self.reduction == "mean": return focal_loss.mean() return focal_loss class HierarchicalLongformer(nn.Module): """CLS-only hierarchical model.""" def __init__(self, config, class_weights_dict): super().__init__() self.config = config self.longformer = LongformerModel.from_pretrained( "allenai/longformer-base-4096", add_pooling_layer=False ) self.longformer.resize_token_embeddings(config.vocab_size) hidden_size = config.hidden_size self.head_binary = nn.Sequential( nn.Dropout(0.1), nn.Linear(hidden_size, hidden_size // 2), nn.LayerNorm(hidden_size // 2), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden_size // 2, 2), ) self.head_family = nn.Sequential( nn.Dropout(0.1), nn.Linear(hidden_size, hidden_size // 2), nn.LayerNorm(hidden_size // 2), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden_size // 2, 4), ) self.head_imp_action = nn.Sequential( nn.Dropout(0.1), nn.Linear(hidden_size, hidden_size // 4), nn.LayerNorm(hidden_size // 4), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden_size // 4, 2), ) self.head_opp_action = nn.Sequential( nn.Dropout(0.1), nn.Linear(hidden_size, hidden_size // 4), nn.LayerNorm(hidden_size // 4), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden_size // 4, 2), ) self.loss_binary = FocalLoss(alpha=class_weights_dict["binary"], gamma=1.5) self.loss_family = FocalLoss(alpha=class_weights_dict["family"], gamma=1.5) self.loss_imp_action = FocalLoss(alpha=class_weights_dict["imp_action"], gamma=2.0) self.loss_opp_action = FocalLoss(alpha=class_weights_dict["opp_action"], gamma=1.5) self.loss_weights = {"binary": 1.0, "family": 1.0, "imp_action": 0.5, "opp_action": 0.5} def forward(self, input_ids, attention_mask, **kwargs): outputs = self.longformer(input_ids=input_ids, attention_mask=attention_mask) cls_output = outputs.last_hidden_state[:, 0, :] return { "logits_binary": self.head_binary(cls_output), "logits_family": self.head_family(cls_output), "logits_imp_action": self.head_imp_action(cls_output), "logits_opp_action": self.head_opp_action(cls_output), } # ── Inference Helpers ──────────────────────────────────────────────────────── def get_final_prediction(logits_binary, logits_family, logits_imp, logits_opp): pred_binary = logits_binary.argmax(dim=-1).item() if pred_binary == 1: return 6 pred_family = logits_family.argmax(dim=-1).item() if pred_family == 0: return 0 elif pred_family == 1: return 1 if logits_imp.argmax(dim=-1).item() == 0 else 2 elif pred_family == 2: return 3 if logits_opp.argmax(dim=-1).item() == 0 else 4 else: return 5 def get_deciding_head(pred_label): if pred_label == 6: return "binary" elif pred_label == 0: return "family" elif pred_label in (1, 2): return "imp_action" elif pred_label in (3, 4): return "opp_action" return "family" # ── Global state (loaded on startup) ──────────────────────────────────────── model = None tokenizer = None device = None no_tags = False def load_model(): """Download model from HF Hub and load into memory.""" global model, tokenizer, device, no_tags device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") # Download files from HF Hub cache_dir = "/tmp/model_cache" os.makedirs(cache_dir, exist_ok=True) print("Downloading model from HF Hub...") model_path = hf_hub_download( repo_id=HF_MODEL_REPO, filename=MODEL_FILENAME, cache_dir=cache_dir, token=HF_TOKEN, ) # Download tokenizer files tokenizer_files = [ "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json", "merges.txt", ] tokenizer_dir = os.path.join(cache_dir, "tokenizer") os.makedirs(tokenizer_dir, exist_ok=True) for fname in tokenizer_files: try: src = hf_hub_download( repo_id=HF_MODEL_REPO, filename=f"{TOKENIZER_DIR}/{fname}", cache_dir=cache_dir, token=HF_TOKEN, ) # Copy to tokenizer dir for from_pretrained import shutil shutil.copy2(src, os.path.join(tokenizer_dir, fname)) except Exception: pass tokenizer = LongformerTokenizerFast.from_pretrained(tokenizer_dir) # Check if trained without tags try: hf_hub_download( repo_id=HF_MODEL_REPO, filename="trained_without_tags.flag", cache_dir=cache_dir, token=HF_TOKEN, ) no_tags = True except Exception: no_tags = False # Infer vocab size from checkpoint try: ckpt = torch.load(model_path, map_location="cpu", weights_only=True) except TypeError: ckpt = torch.load(model_path, map_location="cpu") emb_t = ckpt.get("longformer.embeddings.word_embeddings.weight") if emb_t is None: for k, v in ckpt.items(): if "word_embeddings" in k: emb_t = v break model_vocab_size = int(emb_t.shape[0]) if emb_t is not None else len(tokenizer) del ckpt # Default class weights (only needed for model init, not inference) default_weights = { "binary": torch.ones(2), "family": torch.ones(4), "imp_action": torch.ones(2), "opp_action": torch.ones(2), } config = LongformerConfig.from_pretrained("allenai/longformer-base-4096") config.num_labels = 7 config.vocab_size = model_vocab_size model = HierarchicalLongformer(config, default_weights) try: state_dict = torch.load(model_path, map_location=device, weights_only=True) except TypeError: state_dict = torch.load(model_path, map_location=device) model.load_state_dict(state_dict) model = model.to(device) model.eval() print("✓ Model loaded successfully") # ── FastAPI App ────────────────────────────────────────────────────────────── @asynccontextmanager async def lifespan(app: FastAPI): load_model() yield app = FastAPI( title="Hierarchical Longformer — ESG Classification", description="Live inference for hierarchical ESG text classification", version="1.0.0", lifespan=lifespan, ) # Allow all origins for frontend access app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ── Request/Response Models ────────────────────────────────────────────────── class PredictRequest(BaseModel): text: str class PredictResponse(BaseModel): pred_label: str deciding_head: str probabilities: dict head_predictions: dict # ── Endpoints ──────────────────────────────────────────────────────────────── @app.get("/") async def root(): return { "status": "running", "model": HF_MODEL_REPO, "device": str(device), "labels": LABEL_NAMES, } @app.get("/health") async def health(): return {"status": "healthy", "model_loaded": model is not None} @app.post("/predict", response_model=PredictResponse) async def predict(request: PredictRequest): text = request.text.strip() if not text: return PredictResponse( pred_label="None", deciding_head="binary", probabilities={name: 0.0 for name in LABEL_NAMES}, head_predictions={}, ) # Strip label tags if model was trained without them if no_tags: text = strip_label_tags(text) # Tokenize enc = tokenizer( text, padding="max_length", truncation=True, max_length=2048, return_tensors="pt", ) input_ids = enc["input_ids"].to(device) attention_mask = enc["attention_mask"].to(device) # Run inference with torch.no_grad(): outputs = model(input_ids=input_ids, attention_mask=attention_mask) # Get probabilities probs_binary = F.softmax(outputs["logits_binary"][0], dim=-1).cpu().numpy() probs_family = F.softmax(outputs["logits_family"][0], dim=-1).cpu().numpy() probs_imp = F.softmax(outputs["logits_imp_action"][0], dim=-1).cpu().numpy() probs_opp = F.softmax(outputs["logits_opp_action"][0], dim=-1).cpu().numpy() # Get final prediction pred_label = get_final_prediction( outputs["logits_binary"][0], outputs["logits_family"][0], outputs["logits_imp_action"][0], outputs["logits_opp_action"][0], ) # Compute 7-class probabilities p_esg = float(probs_binary[0]) p_none = float(probs_binary[1]) probabilities = { "Exc": round(p_esg * float(probs_family[0]), 6), "Imp": round(p_esg * float(probs_family[1]) * float(probs_imp[0]), 6), "Imp Act": round(p_esg * float(probs_family[1]) * float(probs_imp[1]), 6), "Opp": round(p_esg * float(probs_family[2]) * float(probs_opp[0]), 6), "Opp Act": round(p_esg * float(probs_family[2]) * float(probs_opp[1]), 6), "Men": round(p_esg * float(probs_family[3]), 6), "None": round(p_none, 6), } # Per-head predictions for transparency head_predictions = { "binary": {"ESG": round(float(probs_binary[0]), 4), "None": round(float(probs_binary[1]), 4)}, "family": { "Exc": round(float(probs_family[0]), 4), "Imp": round(float(probs_family[1]), 4), "Opp": round(float(probs_family[2]), 4), "Men": round(float(probs_family[3]), 4), }, "imp_action": {"Imp": round(float(probs_imp[0]), 4), "Imp Act": round(float(probs_imp[1]), 4)}, "opp_action": {"Opp": round(float(probs_opp[0]), 4), "Opp Act": round(float(probs_opp[1]), 4)}, } return PredictResponse( pred_label=LABEL_NAMES[pred_label], deciding_head=get_deciding_head(pred_label), probabilities=probabilities, head_predictions=head_predictions, )