Miyaka/alertify-system2
0
1"""Alertify AI predictor.2 3This module loads the *real* multi-task mBERT model exported by your Colab trainer:4 5 model/6 tokenizer.json7 tokenizer_config.json8 label_maps.json9 location_gazetteer.json10 heads.ckpt11 encoder/12 config.json13 model.safetensors (optional but recommended)14 15If encoder weights are missing, we gracefully fall back to downloading the16base encoder specified in label_maps.json (Render can do this during build).17"""18 19from __future__ import annotations20 21import json22import os23import re24import unicodedata25from dataclasses import dataclass26from pathlib import Path27from typing import Any, Dict, Tuple28 29import numpy as np30import torch31import torch.nn as nn32from transformers import AutoModel, AutoTokenizer33 34 35def _strip_accents(s: str) -> str:36 s = unicodedata.normalize("NFKD", s)37 return "".join(ch for ch in s if not unicodedata.combining(ch))38 39 40def _squash_repeats(s: str) -> str:41 return re.sub(r"(.)\1{3,}", r"\1\1", s)42 43 44_ROMAN_MAP = {45 "i": "1",46 "ii": "2",47 "iii": "3",48 "iv": "4",49 "v": "5",50 "vi": "6",51 "vii": "7",52 "viii": "8",53 "ix": "9",54 "x": "10",55 "xi": "11",56 "xii": "12",57}58 59_NUM_WORDS = {60 "uno": "1",61 "una": "1",62 "isa": "1",63 "one": "1",64 "wan": "1",65 "dos": "2",66 "dalawa": "2",67 "two": "2",68 "tu": "2",69 "to": "2",70 "tres": "3",71 "tatlo": "3",72 "three": "3",73 "tri": "3",74 "tree": "3",75 "kwatro": "4",76 "cuatro": "4",77 "quatro": "4",78 "apat": "4",79 "four": "4",80 "for": "4",81 "por": "4",82 "singko": "5",83 "cinco": "5",84 "lima": "5",85 "five": "5",86 "fayb": "5",87 "payb": "5",88 "sais": "6",89 "seis": "6",90 "anim": "6",91 "six": "6",92 "siks": "6",93 "siyete": "7",94 "siete": "7",95 "pito": "7",96 "seven": "7",97 "sewen": "7",98 "otso": "8",99 "ocho": "8",100 "walo": "8",101 "eight": "8",102 "eyt": "8",103 "nwebe": "9",104 "nueve": "9",105 "siyam": "9",106 "nine": "9",107 "nayn": "9",108 "diyes": "10",109 "diez": "10",110 "sampu": "10",111 "ten": "10",112 "onse": "11",113 "once": "11",114 "eleven": "11",115 "ileven": "11",116 "dose": "12",117 "doce": "12",118 "twelve": "12",119 "twelv": "12",120}121 122 123def _norm(s: str) -> str:124 s = _strip_accents(str(s).lower()).replace("ñ", "n")125 s = _squash_repeats(s)126 s = re.sub(r"https?://\S+|www\.\S+", " ", s)127 s = re.sub(r"([a-z])\.([a-z])", r"\1 \2", s)128 s = re.sub(r"\bsa([a-z])", r"sa \1", s)129 s = re.sub(r"[^a-z0-9\s\-\/&\.]", " ", s)130 s = re.sub(r"\s+", " ", s).strip()131 132 toks = []133 for t in s.split():134 t0 = t.strip("-")135 toks.append(_ROMAN_MAP.get(t0, t))136 s = " ".join(toks)137 138 toks = [_NUM_WORDS.get(t, t) for t in s.split()]139 s = " ".join(toks)140 141 s = s.replace("&", " and ").replace("/", " ")142 s = re.sub(r"\s+", " ", s).strip()143 return s144 145 146def _softmax(x: np.ndarray) -> np.ndarray:147 x = x - x.max(axis=1, keepdims=True)148 e = np.exp(x)149 return e / (e.sum(axis=1, keepdims=True) + 1e-12)150 151 152class Alertify3Task(nn.Module):153 """Same head architecture used in training (encoder + 3 linear heads)."""154 155 def __init__(self, base_encoder: nn.Module, n_type: int, n_urg: int, n_loc: int, dropout_p: float = 0.15):156 super().__init__()157 self.encoder = base_encoder158 hid = self.encoder.config.hidden_size159 self.dropout = nn.Dropout(dropout_p)160 self.type_head = nn.Linear(hid, n_type)161 self.urg_head = nn.Linear(hid, n_urg)162 self.loc_head = nn.Linear(hid, n_loc)163 164 def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, **kwargs):165 out = self.encoder(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)166 cls = self.dropout(out.last_hidden_state[:, 0])167 lt = self.type_head(cls)168 lu = self.urg_head(cls)169 ll = self.loc_head(cls)170 logits = torch.cat([lt, lu, ll], dim=1)171 return {"logits": logits}172 173 174@dataclass175class Gazetteer:176 barangays: list177 landmarks: list178 alias_index: Dict[str, Dict[str, str]] # norm_alias -> {type, canon}179 180 181class AlertifyAIPredictor:182 """Production-safe predictor used by Flask endpoints."""183 184 def __init__(self, model_dir: str | Path = "model", device: str | None = None):185 self.model_dir = Path(model_dir)186 self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")187 188 # If we can't load the transformer encoder (missing weights and no189 # internet during local dev), we keep the API alive using a190 # lightweight keyword fallback. This prevents the whole web app191 # from crashing, while still allowing deployment with the real model.192 self._disabled = False193 194 self._load_label_maps()195 self._load_tokenizer()196 self._load_gazetteer()197 self._load_model()198 199 # ---------------------------200 # Loaders201 # ---------------------------202 def _load_label_maps(self):203 p = self.model_dir / "label_maps.json"204 if not p.exists():205 raise FileNotFoundError(f"Missing {p}. Put your exported Colab model files under {self.model_dir}.")206 with open(p, "r", encoding="utf-8") as f:207 lm = json.load(f)208 self.base_model = lm.get("base_model", "bert-base-multilingual-cased")209 self.max_len = int(lm.get("max_len", 128))210 self.type_labels = lm["type_labels"]211 self.urg_labels = lm["urg_labels"]212 self.loc_labels = lm["loc_labels"]213 self.id2type = {int(k): v for k, v in lm.get("id2type", {}).items()} if isinstance(lm.get("id2type"), dict) else {i: t for i, t in enumerate(self.type_labels)}214 self.id2urg = {int(k): v for k, v in lm.get("id2urg", {}).items()} if isinstance(lm.get("id2urg"), dict) else {i: t for i, t in enumerate(self.urg_labels)}215 self.id2loc = {int(k): v for k, v in lm.get("id2loc", {}).items()} if isinstance(lm.get("id2loc"), dict) else {i: t for i, t in enumerate(self.loc_labels)}216 self.loc2id = lm.get("loc2id") or {t: i for i, t in enumerate(self.loc_labels)}217 self.none_loc_id = int(self.loc2id.get("NONE", 0))218 219 def _load_tokenizer(self):220 # Prefer local tokenizer files (exported by Colab)221 self.tokenizer = AutoTokenizer.from_pretrained(str(self.model_dir), local_files_only=True)222 223 def _load_gazetteer(self):224 p = self.model_dir / "location_gazetteer.json"225 if not p.exists():226 # still allow running (location will rely on ML only)227 self.gaz = Gazetteer(barangays=[], landmarks=[], alias_index={})228 self._aliases_sorted = []229 return230 with open(p, "r", encoding="utf-8") as f:231 gz = json.load(f)232 self.gaz = Gazetteer(233 barangays=gz.get("barangays", []),234 landmarks=gz.get("landmarks", []),235 alias_index=gz.get("alias_index", {}),236 )237 self._aliases_sorted = sorted(self.gaz.alias_index.keys(), key=len, reverse=True)238 239 self._brgy_re = re.compile(240 r"\b(?:brgy|brgy\.|bgy|bg|bry|barangay|poblacion|pob|pob\.)\s*([a-z0-9\- ]{1,40})",241 re.IGNORECASE,242 )243 self._prep_re = re.compile(244 r"\b(?:sa|nasa|dito\s*sa|d2\s*sa|bandang|malapit\s*sa|near|tapat\s*ng|tapat\s*sa|galing|papuntang|from|going\s*to)\s+([a-z0-9\- ]{2,60})",245 re.IGNORECASE,246 )247 248 def _load_model(self):249 encoder_dir = self.model_dir / "encoder"250 # 1) Prefer local encoder (loads config + weights)251 encoder = None252 if encoder_dir.exists():253 try:254 encoder = AutoModel.from_pretrained(str(encoder_dir), local_files_only=True)255 except Exception:256 encoder = None257 # 2) Fallback to downloading base model (Render build-time).258 # If this environment has no internet access, we degrade gracefully.259 if encoder is None:260 try:261 encoder = AutoModel.from_pretrained(self.base_model)262 except Exception as e:263 # No weights locally and can't download: keep server alive.264 self._disabled = True265 self._disabled_reason = f"encoder_load_failed: {type(e).__name__}" # for debugging266 self.encoder = None267 self.model = None268 return269 270 self.model = Alertify3Task(271 base_encoder=encoder,272 n_type=len(self.type_labels),273 n_urg=len(self.urg_labels),274 n_loc=len(self.loc_labels),275 dropout_p=0.15,276 )277 278 ckpt_path = self.model_dir / "heads.ckpt"279 if not ckpt_path.exists():280 raise FileNotFoundError(f"Missing {ckpt_path}. Put heads.ckpt in {self.model_dir}.")281 ckpt = torch.load(str(ckpt_path), map_location="cpu")282 self.model.type_head.load_state_dict(ckpt["type_head"], strict=True)283 self.model.urg_head.load_state_dict(ckpt["urg_head"], strict=True)284 self.model.loc_head.load_state_dict(ckpt["loc_head"], strict=True)285 self.model.to(self.device)286 self.model.eval()287 288 # cached head segment sizes289 self._a = len(self.type_labels)290 self._b = self._a + len(self.urg_labels)291 292 # ---------------------------293 # Location extraction (gazetteer + simple fuzzy fallback)294 # ---------------------------295 def extract_location(self, text: str) -> Tuple[str, str, int, Dict[str, Any]]:296 t = _norm(text)297 if not t:298 return ("NONE", "", 0, {"reason": "empty"})299 300 # 1) direct alias substring (fast + accurate)301 for a in self._aliases_sorted:302 if len(a) >= 3 and a in t:303 info = self.gaz.alias_index[a]304 conf = 98 if len(a) >= 6 else 92305 return (info.get("type", "NONE"), info.get("canon", ""), conf, {"method": "alias_substring", "alias": a})306 307 # 2) heuristic capture (no rapidfuzz dependency required)308 # If you install rapidfuzz, you can upgrade this further.309 m = self._brgy_re.search(t) if hasattr(self, "_brgy_re") else None310 if m:311 cap = _norm(m.group(1))312 cap = " ".join(cap.split()[:6])313 # cheap similarity: pick alias with longest common substring via contains314 best = None315 best_len = 0316 for a in self._aliases_sorted[:4000]: # cap for speed317 if cap and cap in a and len(cap) > best_len:318 best, best_len = a, len(cap)319 if best:320 info = self.gaz.alias_index[best]321 return (info.get("type", "NONE"), info.get("canon", ""), 86, {"method": "brgy_regex_contains", "captured": cap, "alias": best})322 323 m2 = self._prep_re.search(t) if hasattr(self, "_prep_re") else None324 if m2:325 cap = _norm(m2.group(1))326 cap = " ".join(cap.split()[:8])327 best = None328 best_len = 0329 for a in self._aliases_sorted[:4000]:330 if cap and cap in a and len(cap) > best_len:331 best, best_len = a, len(cap)332 if best:333 info = self.gaz.alias_index[best]334 return (info.get("type", "NONE"), info.get("canon", ""), 86, {"method": "prep_regex_contains", "captured": cap, "alias": best})335 336 return ("NONE", "", 0, {"reason": "no_match", "norm": t})337 338 def _heuristic_predict(self, post_text: str) -> Dict[str, Any]:339 """Fallback prediction when the neural model isn't available.340 341 Uses gazetteer location extraction plus a few keyword rules.342 """343 344 text = str(post_text)345 lt, ll, conf, dbg = self.extract_location(text)346 347 t = _norm(text)348 # very small keyword map (extend anytime)349 DISASTER = [350 ("flood", ["baha", "lubog", "stranded", "tubig"]),351 ("fire", ["sunog", "nasusunog", "bumbero", "usok"]),352 ("earthquake", ["lindol", "uga", "yanig"]),353 ("landslide", ["guho", "landslide", "gumuhit", "gumuho"]),354 ("typhoon/wind", ["bagyo", "hangin", "yero", "signal"]),355 ("power outage", ["brownout", "kuryente", "blackout", "transformer"]),356 ("road blocked", ["barado", "blocked", "sarado", "puno", "nahulog"]),357 ("medical emergency", ["duguan", "nahimatay", "aksidente", "ambulansya"]),358 ("missing person", ["missing", "nawawala", "hinahanap", "last seen"]),359 ]360 361 dtype = "Non-disaster"362 for name, kws in DISASTER:363 if any(k in t for k in kws):364 # Best-effort: choose the closest label if your trained labels exist365 # otherwise keep this pretty name.366 dtype = self.type_labels[0] if self.type_labels else name367 # try exact match368 for lab in self.type_labels:369 if lab.lower() == name.lower():370 dtype = lab371 break372 break373 374 # urgency heuristic375 urgent_tokens = ["tulong", "asap", "now", "agad", "rescue", "trapped", "na-trap", "saklolo", "emergency"]376 critical_tokens = ["mamamatay", "di kami makalabas", "baha hanggang dibdib", "life", "critical"]377 urg = "NON-URGENT"378 if any(k in t for k in urgent_tokens):379 urg = "HIGH"380 if any(k in t for k in critical_tokens):381 urg = "CRITICAL"382 # map to your label set if present383 if self.urg_labels:384 if urg not in self.urg_labels:385 # choose a reasonable fallback386 if "HIGH" in self.urg_labels and urg in ("HIGH", "CRITICAL"):387 urg = "HIGH"388 elif "MODERATE" in self.urg_labels:389 urg = "MODERATE"390 else:391 urg = self.urg_labels[0]392 393 # location decision: gazetteer only394 loc_type, loc_label, loc_src = ("NONE", "", "NONE")395 if lt != "NONE":396 loc_type, loc_label, loc_src = lt, ll, f"GAZETTEER({conf})"397 398 return {399 "input_text": post_text,400 "location_type": loc_type,401 "location_label": loc_label,402 "location_source": loc_src,403 "gazetteer_debug": dbg,404 "disaster_type": dtype,405 "urgency_level": urg,406 "ml_location_top": "NONE",407 "ml_location_prob": 0.0,408 }409 410 # ---------------------------411 # Prediction412 # ---------------------------413 @torch.no_grad()414 def predict(self, post_text: str, use_gazetteer_first: bool = True, loc_prob_threshold: float = 0.45) -> Dict[str, Any]:415 # Fallback mode: if the transformer encoder wasn't loaded (e.g., missing416 # weights and no internet), keep the app functional.417 if getattr(self, "_disabled", False) or self.model is None:418 return self._heuristic_predict(post_text)419 420 g_typ, g_label, g_conf, g_dbg = self.extract_location(post_text)421 422 enc = self.tokenizer(423 post_text,424 truncation=True,425 padding="max_length",426 max_length=self.max_len,427 return_tensors="pt",428 )429 enc = {k: v.to(self.device) for k, v in enc.items()}430 out = self.model(**enc)431 logits = out["logits"].detach().cpu().numpy()432 433 logits_type = logits[:, : self._a]434 logits_urg = logits[:, self._a : self._b]435 logits_loc = logits[:, self._b :]436 437 p_type = int(np.argmax(logits_type, axis=1)[0])438 p_urg = int(np.argmax(logits_urg, axis=1)[0])439 p_loc = int(np.argmax(logits_loc, axis=1)[0])440 441 loc_probs = _softmax(logits_loc)442 loc_prob = float(loc_probs[0, p_loc])443 ml_loc_label = self.id2loc.get(p_loc, "NONE")444 445 chosen_type, chosen_label = "NONE", ""446 chosen_src = "NONE"447 448 # Gazetteer first, else ML with probability threshold449 if use_gazetteer_first and g_typ != "NONE":450 chosen_type, chosen_label, chosen_src = g_typ, g_label, f"GAZETTEER({g_conf})"451 else:452 if ml_loc_label != "NONE" and loc_prob >= loc_prob_threshold:453 if ml_loc_label in self.gaz.barangays:454 chosen_type = "BARANGAY"455 elif ml_loc_label in self.gaz.landmarks:456 chosen_type = "LANDMARK"457 else:458 chosen_type = "LANDMARK"459 chosen_label, chosen_src = ml_loc_label, f"MODEL(p={loc_prob:.2f})"460 elif g_typ != "NONE":461 chosen_type, chosen_label, chosen_src = g_typ, g_label, f"GAZETTEER({g_conf})"462 463 return {464 "input_text": post_text,465 "location_type": chosen_type,466 "location_label": chosen_label,467 "location_source": chosen_src,468 "gazetteer_debug": g_dbg,469 "disaster_type": self.id2type.get(p_type, self.type_labels[p_type]),470 "urgency_level": self.id2urg.get(p_urg, self.urg_labels[p_urg]),471 "ml_location_top": ml_loc_label,472 "ml_location_prob": loc_prob,473 }474