app.py

"""
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,
    )