Spaces:
Running
Running
File size: 3,467 Bytes
4bba8df 8869f68 4bba8df 8869f68 471013f 8869f68 4bba8df af93ecc 4bba8df 8869f68 4bba8df 89d4ec8 8869f68 89d4ec8 4bba8df 8869f68 4bba8df 444ac6d 1849f87 4bba8df 8869f68 4bba8df fb1a253 4bba8df fb1a253 8869f68 fb1a253 4bba8df fb1a253 8869f68 fb1a253 8869f68 853f29a 4bba8df 8869f68 4bba8df 8869f68 4bba8df 8869f68 4bba8df 8869f68 4bba8df 8869f68 4bba8df 8869f68 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import gradio as gr
import os
import torch
import numpy as np
import pandas as pd
from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer
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",
}
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}"
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):
if domain in ["social"]:
return "poltextlab/xlm-roberta-large-twitter-cap-minor"
return "poltextlab/xlm-roberta-large-pooled-cap-minor-v3"
def predict(text, model_id, tokenizer_id):
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()
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
def predict_cap(text, language, domain):
domain = domains[domain]
model_id = build_huggingface_path(language, domain)
tokenizer_id = "xlm-roberta-large"
if is_disk_full():
os.system("rm -rf /data/models*")
os.system("rm -r ~/.cache/huggingface/hub")
return predict(text, model_id, tokenizer_id)
demo = gr.Interface(
title="CAP Minor Topics Babel Demo",
fn=predict_cap,
inputs=[
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.Label(num_top_classes=5, label="Output"), gr.Markdown()],
)
|