EuroVoc / predict.py

predict.py: skip content-free input (empty/zero-width/emoji) to prevent unrelated tags on empty input (#3)

eea3ce4 about 1 month ago

2.07 kB

"""Predict EuroVoc tags from the HF layout (encoder model.safetensors + classifier.pt).
Shipped into the HF repo as predict.py. Run from the model directory:
    uv run python predict.py "some EU document text"     # or no arg -> demo doc
Thresholds are the tuned ship operating points, not 0.5 (see MODEL_CARD.md).
"""
import json
import pickle
import sys
from pathlib import Path

import torch
from transformers import AutoModel, AutoTokenizer

HERE = Path(__file__).parent
CONCEPT_THRESHOLD = 0.9          # concept level; coarser levels use MT 0.80, DO 0.70


def load():
    tok = AutoTokenizer.from_pretrained(HERE)
    enc = AutoModel.from_pretrained(HERE).eval()
    head = torch.load(HERE / "classifier.pt", weights_only=True)
    mlb = pickle.load(open(HERE / "mlb.pickle", "rb"))
    names = json.loads((HERE / "names.json").read_text())
    return tok, enc, head, mlb, names


def predict(text, thresh=CONCEPT_THRESHOLD, model=None):
    if not any(c.isalnum() for c in text):   # no letters/digits in any script: only special tokens would
        return []                            # be pooled, which yields hundreds of spurious tags. Skip it.
    tok, enc, head, mlb, names = model or load()   # reloads per call; pass model= in a loop
    b = tok(text, truncation=True, max_length=2048, return_tensors="pt")
    with torch.no_grad():
        h = enc(**b).last_hidden_state
        m = b["attention_mask"].unsqueeze(-1).float()
        p = ((h * m).sum(1) / m.sum(1).clamp(min=1e-9) @ head["weight"].T + head["bias"]).sigmoid()[0]
    tags = [(names.get(str(mlb.classes_[i]), str(mlb.classes_[i])), float(p[i]))
            for i in (p > thresh).nonzero().flatten().tolist()]
    return sorted(tags, key=lambda t: -t[1])


if __name__ == "__main__":
    demo = "The European Parliament adopted a regulation on renewable energy and climate policy."
    text = sys.argv[1] if len(sys.argv) > 1 else demo
    hits = predict(text)
    print(f"{len(hits)} tags over {CONCEPT_THRESHOLD}:")
    for name, score in hits:
        print(f"  {score:.2f}  {name}")

	"""Predict EuroVoc tags from the HF layout (encoder model.safetensors + classifier.pt).
	Shipped into the HF repo as predict.py. Run from the model directory:

	uv run python predict.py "some EU document text" # or no arg -> demo doc

	Thresholds are the tuned ship operating points, not 0.5 (see MODEL_CARD.md).
	"""
	import json
	import pickle
	import sys
	from pathlib import Path

	import torch
	from transformers import AutoModel, AutoTokenizer

	HERE = Path(__file__).parent
	CONCEPT_THRESHOLD = 0.9 # concept level; coarser levels use MT 0.80, DO 0.70


	def load():
	tok = AutoTokenizer.from_pretrained(HERE)
	enc = AutoModel.from_pretrained(HERE).eval()
	head = torch.load(HERE / "classifier.pt", weights_only=True)
	mlb = pickle.load(open(HERE / "mlb.pickle", "rb"))
	names = json.loads((HERE / "names.json").read_text())
	return tok, enc, head, mlb, names


	def predict(text, thresh=CONCEPT_THRESHOLD, model=None):
	if not any(c.isalnum() for c in text): # no letters/digits in any script: only special tokens would
	return [] # be pooled, which yields hundreds of spurious tags. Skip it.
	tok, enc, head, mlb, names = model or load() # reloads per call; pass model= in a loop
	b = tok(text, truncation=True, max_length=2048, return_tensors="pt")
	with torch.no_grad():
	h = enc(**b).last_hidden_state
	m = b["attention_mask"].unsqueeze(-1).float()
	p = ((h * m).sum(1) / m.sum(1).clamp(min=1e-9) @ head["weight"].T + head["bias"]).sigmoid()[0]
	tags = [(names.get(str(mlb.classes_[i]), str(mlb.classes_[i])), float(p[i]))
	for i in (p > thresh).nonzero().flatten().tolist()]
	return sorted(tags, key=lambda t: -t[1])


	if __name__ == "__main__":
	demo = "The European Parliament adopted a regulation on renewable energy and climate policy."
	text = sys.argv[1] if len(sys.argv) > 1 else demo
	hits = predict(text)
	print(f"{len(hits)} tags over {CONCEPT_THRESHOLD}:")
	for name, score in hits:
	print(f" {score:.2f} {name}")