Upload 6-head multi-task classifier (peril+severity+category+fire_sub+relevance+actionability)

handler.py

@@ -6,9 +6,10 @@ PERIL_LABELS = ["fire","flood","named_windstorm","construction_theft","transient
 SEVERITY_LABELS = ["low","medium","high","critical"]
 CATEGORY_LABELS = ["incident_report","trend","regulatory","research","warning"]
 FIRE_SUBCATEGORY_LABELS = ["arson","wildfire","unknown_cause"]
+ACTIONABILITY_LABELS = ["irrelevant","informational","notable","actionable"]
 
 class MultiTaskClassifier(nn.Module):
-    def __init__(self, model_name, np, ns, nc, nf):
+    def __init__(self, model_name, np, ns, nc, nf, na=4):
         super().__init__()
         self.encoder = AutoModel.from_pretrained(model_name)
         h = self.encoder.config.hidden_size
@@ -17,17 +18,20 @@ class MultiTaskClassifier(nn.Module):
         self.severity_head = nn.Linear(h, ns)
         self.category_head = nn.Linear(h, nc)
         self.fire_sub_head = nn.Linear(h, nf)
+        self.relevance_head = nn.Linear(h, 1)
+        self.actionability_head = nn.Linear(h, na)
     def forward(self, input_ids, attention_mask=None):
         out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
         pooled = self.dropout(out.last_hidden_state[:, 0, :])
         return {"peril_logits": self.peril_head(pooled), "severity_logits": self.severity_head(pooled),
-                "category_logits": self.category_head(pooled), "fire_sub_logits": self.fire_sub_head(pooled)}
+                "category_logits": self.category_head(pooled), "fire_sub_logits": self.fire_sub_head(pooled),
+                "relevance_logits": self.relevance_head(pooled), "actionability_logits": self.actionability_head(pooled)}
 
 class EndpointHandler:
     def __init__(self, path=""):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = MultiTaskClassifier(path, len(PERIL_LABELS), len(SEVERITY_LABELS), len(CATEGORY_LABELS), len(FIRE_SUBCATEGORY_LABELS))
+        self.model = MultiTaskClassifier(path, len(PERIL_LABELS), len(SEVERITY_LABELS), len(CATEGORY_LABELS), len(FIRE_SUBCATEGORY_LABELS), len(ACTIONABILITY_LABELS))
         w = Path(path) / "pytorch_model.bin"
         if w.exists():
             self.model.load_state_dict(torch.load(str(w), map_location=self.device, weights_only=True))
@@ -36,6 +40,9 @@ class EndpointHandler:
     def __call__(self, data):
         text = data.get("inputs", "")
         if isinstance(text, list): text = text[0]
+        params = data.get("parameters", {})
+        include_relevance = params.get("include_relevance", False)
+        include_actionability = params.get("include_actionability", False)
         inputs = self.tokenizer(text[:16000], truncation=True, max_length=512, return_tensors="pt").to(self.device)
         with torch.no_grad():
             out = self.model(**inputs)
@@ -46,7 +53,15 @@ class EndpointHandler:
         ci = int(out["category_logits"].argmax(-1)[0].cpu())
         fp = torch.softmax(out["fire_sub_logits"], -1)[0].cpu().tolist()
         fi = int(out["fire_sub_logits"].argmax(-1)[0].cpu())
-        return {"peril_scores": {l: round(s,4) for l,s in zip(PERIL_LABELS, pp)},
-                "severity": {"label": SEVERITY_LABELS[si], "confidence": round(sp[si],4)},
-                "category": {"label": CATEGORY_LABELS[ci], "confidence": round(cp[ci],4)},
-                "fire_subcategory": {"label": FIRE_SUBCATEGORY_LABELS[fi], "confidence": round(fp[fi],4)}}
+        result = {"peril_scores": {l: round(s,4) for l,s in zip(PERIL_LABELS, pp)},
+                  "severity": {"label": SEVERITY_LABELS[si], "confidence": round(sp[si],4)},
+                  "category": {"label": CATEGORY_LABELS[ci], "confidence": round(cp[ci],4)},
+                  "fire_subcategory": {"label": FIRE_SUBCATEGORY_LABELS[fi], "confidence": round(fp[fi],4)}}
+        if include_relevance:
+            rel_score = float(torch.sigmoid(out["relevance_logits"].squeeze(-1))[0].cpu())
+            result["relevance"] = {"score": round(rel_score, 4), "label": "relevant" if rel_score > 0.5 else "irrelevant"}
+        if include_actionability:
+            ap = torch.softmax(out["actionability_logits"], -1)[0].cpu().tolist()
+            ai = int(out["actionability_logits"].argmax(-1)[0].cpu())
+            result["actionability"] = {"label": ACTIONABILITY_LABELS[ai], "confidence": round(ap[ai], 4)}
+        return result

label_config.json

@@ -25,5 +25,11 @@
     "0": "arson",
     "1": "wildfire",
     "2": "unknown_cause"
+  },
+  "actionability_labels": {
+    "0": "irrelevant",
+    "1": "informational",
+    "2": "notable",
+    "3": "actionable"
   }
 }

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2bd721a7ff60be333ebd50cc28b86a2eab55037ab3fdc135e04150a778ee57f7
-size 498727455
+oid sha256:33da6dbfb4dc60916ae56804d7fe4affe89b7fd456fc89aa652cdc7e34c6cdf8
+size 498744175