Spaces:
Running
Running
import gradio as gr | |
import os | |
import torch | |
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 ( | |
CAP_NUM_DICT, | |
CAP_LABEL_NAMES, | |
CAP_MIN_NUM_DICT, | |
CAP_MIN_LABEL_NAMES, | |
) | |
from .utils import is_disk_full, release_model | |
HF_TOKEN = os.environ["hf_read"] | |
languages = [ | |
"Multilingual", | |
] | |
domains = { | |
"media": "media", | |
"social media": "social", | |
"parliamentary speech": "parlspeech", | |
"legislative documents": "legislative", | |
"executive speech": "execspeech", | |
"executive order": "execorder", | |
"party programs": "party", | |
"judiciary": "judiciary", | |
"budget": "budget", | |
"public opinion": "publicopinion", | |
"local government agenda": "localgovernment", | |
} | |
CAP_MAJOR_CODES = list(CAP_NUM_DICT.values()) | |
CAP_MIN_CODES = list(CAP_MIN_NUM_DICT.values()) | |
major_index_to_id = {i: code for i, code in enumerate(CAP_MAJOR_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): | |
try: | |
hf_api = HfApi(token=HF_TOKEN) | |
hf_api.model_info(checkpoint_path, token=HF_TOKEN) | |
return True | |
except: | |
return False | |
def build_huggingface_path(language: str, domain: str, hierarchical: bool): | |
if hierarchical: | |
major = "poltextlab/xlm-roberta-large-pooled-cap-v3" | |
minor = ( | |
"poltextlab/xlm-roberta-large-twitter-cap-minor" | |
if domain == "social" | |
else "poltextlab/xlm-roberta-large-pooled-cap-minor-v3" | |
) | |
return major, minor | |
return ( | |
"poltextlab/xlm-roberta-large-twitter-cap-minor" | |
if domain == "social" | |
else "poltextlab/xlm-roberta-large-pooled-cap-minor-v3" | |
) | |
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 | |
if int(top_major_id) == 999: | |
filtered_probs = {21:1} | |
else: | |
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_LABEL_NAMES[major_index_to_id[k]]}": v | |
for k, v in sorted( | |
filtered_probs.items(), key=lambda item: item[1], reverse=True | |
) | |
} | |
print(top_major_id) # debug | |
print(major_to_minor_map) # debug | |
# 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_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 | |
jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt" | |
model = torch.jit.load(jit_model_path).to(device) | |
model.eval() | |
# Load tokenizer (still regular HF) | |
tokenizer = AutoTokenizer.from_pretrained(tokenizer_id) | |
# Tokenize input | |
inputs = tokenizer( | |
text, max_length=64, truncation=True, padding=True, return_tensors="pt" | |
) | |
inputs = {k: v.to(device) for k, v in inputs.items()} | |
with torch.no_grad(): | |
output = model(inputs["input_ids"], inputs["attention_mask"]) | |
print(output) # debug | |
logits = output["logits"] | |
release_model(model, model_id) | |
probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten() | |
top_indices = np.argsort(probs)[::-1][:10] | |
output_pred = {} | |
for i in top_indices: | |
code = CAP_MIN_NUM_DICT[i] | |
prob = probs[i] | |
if code in CAP_LABEL_NAMES: | |
# Media (major) topic | |
label = CAP_LABEL_NAMES[code] | |
display = f"[{code}] {label}" | |
else: | |
# Minor topic | |
major_code = code // 100 | |
major_label = CAP_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_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 interpretation_info, output_pred, output_info | |
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") | |
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]), | |
], | |
outputs=[gr.Markdown(), gr.Label(label="Output"), gr.Markdown()], | |
) | |