File size: 4,131 Bytes
567677e c9193b2 567677e c9193b2 567677e 67d49c1 567677e 67d49c1 567677e 67d49c1 567677e 67d49c1 567677e 67d49c1 567677e | 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 | from typing import Any, Dict, List
import torch
from transformers import BertTokenizerFast
from model import BertForTokenAndSequenceJointClassification, TOKEN_TAGS, SEQUENCE_TAGS
MAX_LENGTH = 512
# Special/non-technique tags at the front of TOKEN_TAGS (see model.py).
IGNORED_TAG_IDS = {0, 1} # "<PAD>", "O"
# This repo ships weights + config only, no tokenizer files (vocab.txt /
# tokenizer_config.json). config.json's vocab_size (28996) matches
# bert-base-cased exactly, so the tokenizer is loaded from the public base
# checkpoint instead of the local repo path.
TOKENIZER_BASE = "bert-base-cased"
class EndpointHandler:
def __init__(self, path: str = ""):
self.tokenizer = BertTokenizerFast.from_pretrained(TOKENIZER_BASE)
self.model = BertForTokenAndSequenceJointClassification.from_pretrained(path)
self.model.eval()
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
text = data.get("inputs")
if not isinstance(text, str) or not text:
return {"error": "expected data.inputs to be a non-empty string"}
encoding = self.tokenizer(
text,
return_offsets_mapping=True,
return_tensors="pt",
truncation=True,
max_length=MAX_LENGTH,
)
offset_mapping = encoding.pop("offset_mapping")[0].tolist()
with torch.no_grad():
outputs = self.model(**encoding)
# token_logits: [1, seq_len, 20] at inference (labels=None), see model.py forward().
token_probs = torch.softmax(outputs.token_logits[0], dim=-1)
token_tag_ids = token_probs.argmax(dim=-1).tolist()
token_tag_probs = token_probs.max(dim=-1).values.tolist()
sequence_probs = torch.softmax(outputs.sequence_logits[0], dim=-1).tolist()
sequence_tag_id = int(torch.argmax(outputs.sequence_logits[0]).item())
# Merge consecutive TOKENS (by index, not character offset) sharing
# the same non-ignored tag into one span. PTC/this model has no BIO
# scheme: adjacent same-tag tokens belong to the same span by
# construction (see GUIDA_MIGRAZIONE, §1). Character offsets are NOT
# a valid adjacency test here: a word boundary space makes the next
# token's start > the previous token's end even though they're the
# same run of tokens, so merging must key off token index, not char
# offset. Special tokens ([CLS]/[SEP]/[PAD]) have offset (0, 0) and
# are skipped naturally since they never carry a real technique tag.
spans: List[Dict[str, Any]] = []
current = None
for i, (tag_id, prob, (start, end)) in enumerate(zip(token_tag_ids, token_tag_probs, offset_mapping)):
is_technique = tag_id not in IGNORED_TAG_IDS and end > start
if (
is_technique
and current is not None
and current["tag_id"] == tag_id
and i == current["last_index"] + 1
):
current["end"] = end
current["probs"].append(prob)
current["last_index"] = i
else:
if current is not None:
spans.append(current)
current = (
{"tag_id": tag_id, "start": start, "end": end, "probs": [prob], "last_index": i}
if is_technique
else None
)
if current is not None:
spans.append(current)
out_spans = [
{
"technique": TOKEN_TAGS[s["tag_id"]],
"start": s["start"],
"end": s["end"],
"text": text[s["start"] : s["end"]],
"confidence": sum(s["probs"]) / len(s["probs"]),
}
for s in spans
]
return {
"spans": out_spans,
"sequence_label": SEQUENCE_TAGS[sequence_tag_id],
"sequence_probs": {SEQUENCE_TAGS[i]: p for i, p in enumerate(sequence_probs)},
"truncated": len(encoding["input_ids"][0]) >= MAX_LENGTH,
}
|