1import json
2import numpy as np
3import torch
4from huggingface_hub import hf_hub_download
5from transformers import XLMRobertaTokenizerFast, XLMRobertaForTokenClassification
6
7model_name = "Setur/BRAGD"
8
9tokenizer = XLMRobertaTokenizerFast.from_pretrained(model_name)
10model = XLMRobertaForTokenClassification.from_pretrained(model_name)
11model.eval()
12
13# Download decoding assets
14constraint_mask_path = hf_hub_download(model_name, "constraint_mask.json")
15tag_mappings_path = hf_hub_download(model_name, "tag_mappings.json")
16
17with open(constraint_mask_path, "r", encoding="utf-8") as f:
18 raw_mask = json.load(f)
19constraint_mask = {int(k): [tuple(x) for x in v] for k, v in raw_mask.items()}
20
21with open(tag_mappings_path, "r", encoding="utf-8") as f:
22 raw_map = json.load(f)
23features_to_tag = {tuple(map(int, k.split(","))): v for k, v in raw_map.items()}
24
25WORD_CLASS_NAMES = {
26 0: "Noun",
27 1: "Adjective",
28 2: "Pronoun",
29 3: "Number",
30 4: "Verb",
31 5: "Participle",
32 6: "Adverb",
33 7: "Conjunction",
34 8: "Foreign",
35 9: "Unanalyzed",
36 10: "Abbreviation",
37 11: "Web",
38 12: "Punctuation",
39 13: "Symbol",
40 14: "Article",
41}
42
43INTERVAL_NAMES = {
44 (15, 29): "subcategory",
45 (30, 33): "gender",
46 (34, 36): "number",
47 (37, 41): "case",
48 (42, 43): "article",
49 (44, 45): "proper_noun",
50 (46, 50): "degree",
51 (51, 53): "declension",
52 (54, 60): "mood",
53 (61, 63): "voice",
54 (64, 66): "tense",
55 (67, 70): "person",
56 (71, 72): "definiteness",
57}
58
59FEATURE_COLUMNS = [
60 "S", "A", "P", "N", "V", "L", "D", "C", "F", "X", "T", "W", "K", "M", "R",
61 "D", "B", "E", "I", "P", "Q", "N", "G", "R", "X", "S", "C", "O", "T", "s",
62 "M", "F", "N", "g",
63 "S", "P", "n",
64 "N", "A", "D", "G", "c",
65 "A", "a",
66 "P", "r",
67 "P", "C", "S", "A", "d",
68 "S", "W", "e",
69 "I", "M", "N", "S", "P", "E", "U",
70 "A", "M", "v",
71 "P", "A", "t",
72 "1", "2", "3", "p",
73 "D", "I",
74]
75
76def decode_token(logits):
77 pred = np.zeros(logits.shape[0], dtype=int)
78
79 # predict word class
80 wc = int(np.argmax(logits[:15]))
81 pred[wc] = 1
82
83 # predict only valid feature groups for this word class
84 for start, end in constraint_mask.get(wc, []):
85 group = logits[start:end+1]
86 pred[start + int(np.argmax(group))] = 1
87
88 tag = features_to_tag.get(tuple(pred.tolist()), None)
89
90 features = {"word_class": WORD_CLASS_NAMES.get(wc, str(wc))}
91 for (start, end), name in INTERVAL_NAMES.items():
92 group = pred[start:end+1]
93 active = np.where(group == 1)[0]
94 if len(active) == 1:
95 features[name] = FEATURE_COLUMNS[start + active[0]]
96
97 return tag, features
98
99text = "Hetta er eitt føroyskt dømi"
100words = text.split()
101
102enc = tokenizer(
103 [words],
104 is_split_into_words=True,
105 return_tensors="pt",
106 padding=True,
107 truncation=True,
108)
109
110with torch.no_grad():
111 logits = model(**enc).logits[0]
112
113word_ids = enc.word_ids(batch_index=0)
114seen = set()
115
116for i, word_id in enumerate(word_ids):
117 if word_id is None or word_id in seen:
118 continue
119 seen.add(word_id)
120
121 tag, features = decode_token(logits[i].cpu().numpy())
122 print(f"{words[word_id]:15s} {str(tag):10s} {features}")