cap minor hierarchical

interfaces/cap_minor.py

@@ -6,8 +6,8 @@ import numpy as np
 import pandas as pd
 from transformers import AutoModelForSequenceClassification
 from transformers import AutoTokenizer
+import torch.nn.functional as F
 from huggingface_hub import HfApi
-
 from collections import defaultdict
 
 from label_dicts import (
@@ -41,12 +41,23 @@ domains = {
 }
 
 
-def get_label_name(idx):
-    minor_code = CAP_MIN_NUM_DICT[idx]
-    minor_label_name = CAP_MIN_LABEL_NAMES[minor_code]
-    major_code = minor_code // 100 if minor_code not in [99, 999, 9999] else 999
-    major_label_name = CAP_LABEL_NAMES[major_code]
-    return f"[{major_code}] {major_label_name} [{minor_code}] {minor_label_name}"
+CAP_MEDIA_CODES = list(CAP_MEDIA_NUM_DICT.values())
+CAP_MIN_CODES = list(CAP_MIN_NUM_DICT.values())
+
+major_index_to_id = {i: code for i, code in enumerate(CAP_MEDIA_CODES)}
+minor_id_to_index = {code: i for i, code in enumerate(CAP_MIN_CODES)}
+minor_index_to_id = {i: code for i, code in enumerate(CAP_MIN_CODES)}
+
+major_to_minor_map = defaultdict(list)
+for code in CAP_MIN_CODES:
+    major_id = int(str(code)[:-2])
+    major_to_minor_map[major_id].append(code)
+major_to_minor_map = dict(major_to_minor_map)
+
+
+def normalize_probs(probs: dict) -> dict:
+    total = sum(probs.values())
+    return {k: v / total for k, v in probs.items()}
 
 
 def check_huggingface_path(checkpoint_path: str):
@@ -64,7 +75,99 @@ def build_huggingface_path(language: str, domain: str):
     return "poltextlab/xlm-roberta-large-pooled-cap-minor-v3"
 
 
-def predict(text, model_id, tokenizer_id):
+def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
+    device = torch.device("cpu")
+
+    # Load major and minor models + tokenizer
+    major_model = AutoModelForSequenceClassification.from_pretrained(
+        major_model_id,
+        low_cpu_mem_usage=True,
+        device_map="auto",
+        offload_folder="offload",
+        token=HF_TOKEN,
+    ).to(device)
+
+    minor_model = AutoModelForSequenceClassification.from_pretrained(
+        minor_model_id,
+        low_cpu_mem_usage=True,
+        device_map="auto",
+        offload_folder="offload",
+        token=HF_TOKEN,
+    ).to(device)
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+
+    # Tokenize input
+    inputs = tokenizer(
+        text, max_length=64, truncation=True, padding=True, return_tensors="pt"
+    ).to(device)
+
+    # Predict major topic
+    major_model.eval()
+    with torch.no_grad():
+        major_logits = major_model(**inputs).logits
+        major_probs = F.softmax(major_logits, dim=-1)
+    major_probs_np = major_probs.cpu().numpy().flatten()
+    top_major_index = int(np.argmax(major_probs_np))
+    top_major_id = major_index_to_id[top_major_index]
+
+    # Default: show major topic predictions
+    filtered_probs = {
+        i: float(major_probs_np[i]) for i in np.argsort(major_probs_np)[::-1]
+    }
+    filtered_probs = normalize_probs(filtered_probs)
+
+    output_pred = {
+        f"[{major_index_to_id[k]}] {CAP_MEDIA_LABEL_NAMES[major_index_to_id[k]]}": v
+        for k, v in sorted(
+            filtered_probs.items(), key=lambda item: item[1], reverse=True
+        )
+    }
+
+    # If eligible for minor prediction
+    if top_major_id in major_to_minor_map:
+        valid_minor_ids = major_to_minor_map[top_major_id]
+        minor_model.eval()
+        with torch.no_grad():
+            minor_logits = minor_model(**inputs).logits
+            minor_probs = F.softmax(minor_logits, dim=-1)
+
+        release_model(major_model, major_model_id)
+        release_model(minor_model, minor_model_id)
+
+        print(minor_probs)  # debug
+        # Restrict to valid minor codes
+        valid_indices = [
+            minor_id_to_index[mid]
+            for mid in valid_minor_ids
+            if mid in minor_id_to_index
+        ]
+        filtered_probs = {
+            minor_index_to_id[i]: float(minor_probs[0][i]) for i in valid_indices
+        }
+        print(filtered_probs)  # debug
+        filtered_probs = normalize_probs(filtered_probs)
+        print(filtered_probs)  # debug
+
+        output_pred = {
+            f"[{top_major_id}] {CAP_MEDIA_LABEL_NAMES[top_major_id]} [{k}] {CAP_MIN_LABEL_NAMES[k]}": v
+            for k, v in sorted(
+                filtered_probs.items(), key=lambda item: item[1], reverse=True
+            )
+        }
+
+    output_info = f'<p style="text-align: center; display: block">Prediction used <a href="https://huggingface.co/{major_model_id}">{major_model_id}</a> and <a href="https://huggingface.co/{minor_model_id}">{minor_model_id}</a>.</p>'
+
+    interpretation_info = """
+    ## How to Interpret These Values (Hierarchical Classification)
+
+    The values are the confidences for minor topics **within a given major topic**, and they are **normalized to sum to 1**.
+    """
+
+    return interpretation_info, output_pred, output_info
+
+
+def predict_flat(text, model_id, tokenizer_id, HF_TOKEN=None):
     device = torch.device("cpu")
 
     # Load JIT-traced model
@@ -89,28 +192,77 @@ def predict(text, model_id, tokenizer_id):
     release_model(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
+    top_indices = np.argsort(probs)[::-1][:10]
+
+    CAP_MIN_MEDIA_LABEL_NAMES = CAP_MEDIA_LABEL_NAMES | CAP_MIN_LABEL_NAMES
+
+    output_pred = {}
+    for i in top_indices:
+        code = CAP_MIN_MEDIA_NUM_DICT[i]
+        prob = probs[i]
+
+        if code in CAP_MEDIA_LABEL_NAMES:
+            # Media (major) topic
+            label = CAP_MEDIA_LABEL_NAMES[code]
+            display = f"[{code}] {label}"
+        else:
+            # Minor topic
+            major_code = code // 100
+            major_label = CAP_MEDIA_LABEL_NAMES[major_code]
+            minor_label = CAP_MIN_LABEL_NAMES[code]
+            display = f"[{major_code}] {major_label} [{code}] {minor_label}"
+
+        output_pred[display] = prob
+
+    interpretation_info = """
+    ## How to Interpret These Values (Flat Classification)
+
+    This method returns predictions made by a single model. **Only the top 10 most confident labels are displayed**.
+    """
 
-    output_pred = {get_label_name(i): probs[i] for i in np.argsort(probs)[::-1]}
     output_info = f'<p style="text-align: center; display: block">Prediction was made using the <a href="https://huggingface.co/{model_id}">{model_id}</a> model.</p>'
-    return output_pred, output_info
 
+    return interpretation_info, output_pred, output_info
 
-def predict_cap(text, language, domain):
-    domain = domains[domain]
-    model_id = build_huggingface_path(language, domain)
-    tokenizer_id = "xlm-roberta-large"
 
+def predict_cap(tmp, method, text, language, domain):
     if is_disk_full():
         os.system("rm -rf /data/models*")
         os.system("rm -r ~/.cache/huggingface/hub")
 
-    return predict(text, model_id, tokenizer_id)
+    domain = domains[domain]
+
+    if method == "Hierarchical Classification":
+        major_model_id, minor_model_id = build_huggingface_path(language, domain, True)
+        tokenizer_id = "xlm-roberta-large"
+        return predict(text, major_model_id, minor_model_id, tokenizer_id)
+
+    else:
+        model_id = build_huggingface_path(language, domain, False)
+        tokenizer_id = "xlm-roberta-large"
+        return predict_flat(text, model_id, tokenizer_id)
+
+
+description = """
+You can choose between two approaches for making predictions:
+
+**1. Hierarchical Classification**  
+First, the model predicts a **major topic**. Then, a second model selects the most probable **subtopic** from within that major topic's category.
 
+**2. Flat Classification (single model)**  
+A single model directly predicts the most relevant label from all available minor topics.
+"""
 
 demo = gr.Interface(
     title="CAP Minor Topics Babel Demo",
     fn=predict_cap,
     inputs=[
+        gr.Markdown(description),
+        gr.Radio(
+            choices=["Hierarchical Classification", "Flat Classification"],
+            label="Prediction Mode",
+            value="Hierarchical Classification",
+        ),
         gr.Textbox(lines=6, label="Input"),
         gr.Dropdown(languages, label="Language", value=languages[0]),
         gr.Dropdown(domains.keys(), label="Domain", value=list(domains.keys())[0]),