Views
No views yet

4096 tokens.768-dimensional L2-normalised vector space, with type-aware conditioning to separate coarse types (PER/ORG/LOC).512 tokens, which can be insufficient for documents with multiple entities or complex structures. This model was trained with a context of 4096 tokens.Alibaba-NLP/gte-multilingual-mlm-base) + token-classification head + attention pooling + per entity-type projection head + CRFmention_mask to obtain L2-normalised 768-dimensional span vectors suitable for clustering, retrieval, and linking tasks.1import torch
2import torch.nn.functional as F
3from transformers import AutoModel, AutoTokenizer
4
5model_id = "pierre-tassel/rapido-ner-entity"
6
7model = AutoModel.from_pretrained(model_id, trust_remote_code=True).eval()
8tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
9
10texts = [
11 "I'm traveling to Cologne next week.",
12 "Ich fahre nächste Woche nach Köln.",
13 "Je vais à Cologne la semaine prochaine.",
14 "Apple opened a lab in Paris.",
15 "Microsoft acquired a startup in Berlin.",
16 # one document with multiple mentions to show within-doc clustering (two Cologne LOC)
17 "The Cologne office met with Köln University and Microsoft in Cologne.",
18]
19
20enc = tokenizer(
21 texts,
22 return_tensors="pt",
23 padding=True,
24 truncation=True,
25)
26input_ids = enc["input_ids"]
27attention_mask = enc["attention_mask"].bool()
28
29with torch.no_grad():
30 hidden = model.encode_tokens(input_ids, attention_mask)
31 logits = model.ner_head(hidden)
32
33decoded = model.crf_decode(logits, attention_mask)
34
35id2label = {int(k): v for k, v in model.config.id2label.items()}
36
37
38def spans_from_bio(tag_ids: list[int]) -> list[tuple[int, int, str]]:
39 spans: list[tuple[int, int, str]] = []
40 start: int | None = None
41 cur_type: str | None = None
42
43 for idx, tag_id in enumerate(tag_ids):
44 tag = id2label[tag_id]
45
46 if tag == "O":
47 if start is not None:
48 spans.append((start, idx, cur_type))
49 start, cur_type = None, None
50 continue
51
52 prefix, etype = tag.split("-", 1) if "-" in tag else ("B", tag)
53 if prefix == "B" or etype != cur_type:
54 if start is not None:
55 spans.append((start, idx, cur_type))
56 start, cur_type = idx, etype
57
58 if start is not None:
59 spans.append((start, len(tag_ids), cur_type))
60 return spans
61
62
63results = []
64
65for i, text in enumerate(texts):
66 seq_len = int(attention_mask[i].sum().item())
67 tokens = tokenizer.convert_ids_to_tokens(input_ids[i, :seq_len])
68 tag_ids = decoded[i][:seq_len]
69 spans = spans_from_bio(tag_ids)
70
71 # Print predictions per text
72 tag_labels = [id2label[t] for t in tag_ids]
73 print(f"\nPrediction for text {i + 1}: {text}")
74 print("Tokens with tags:")
75 print(list(zip(tokens, tag_labels))) # e.g., ... ('▁Colo','B-LOC'), ('gne','I-LOC'), ...
76 if spans:
77 ents_readable = [
78 f"{tokenizer.decode(input_ids[i, s:e]).strip()} [{typ}]"
79 for (s, e, typ) in spans
80 ]
81 print("Entities:", ents_readable) # e.g., ['Apple [ORG]', 'Paris [LOC]']
82 else:
83 print("Entities: []")
84
85 mention_vectors = []
86 mention_surfaces = []
87 mention_types = []
88
89 if spans:
90 mention_mask = torch.zeros(
91 (1, len(spans), input_ids.size(1)), dtype=torch.bool
92 )
93 for m, (s, e, _) in enumerate(spans):
94 mention_mask[0, m, s:e] = True
95
96 with torch.no_grad():
97 pooled = model.encode_mentions_with_attention(
98 hidden[i: i + 1], mention_mask
99 )
100 projected = model.project_mentions(pooled).squeeze(0)
101
102 mention_vectors = F.normalize(projected, dim=-1)
103 for (s, e, typ) in spans:
104 surface = tokenizer.decode(input_ids[i, s:e])
105 mention_surfaces.append(surface.strip())
106 mention_types.append(typ)
107
108 results.append(
109 {
110 "text": text,
111 "tokens": tokens,
112 "tags": tag_labels,
113 "spans": spans,
114 "mention_surfaces": mention_surfaces,
115 "mention_types": mention_types,
116 "embeddings": mention_vectors, # tensor or []
117 }
118 )
119
120# Global pairwise similarities
121all_embeds = []
122all_labels = []
123for r in results:
124 for surface, typ, emb in zip(
125 r["mention_surfaces"], r["mention_types"], r["embeddings"]
126 ):
127 all_embeds.append(emb)
128 all_labels.append(f"{surface} [{typ}]")
129
130if all_embeds:
131 all_embeds = torch.stack(all_embeds)
132 sim = all_embeds @ all_embeds.T
133 print(
134 "\nGlobal cosine similarities:" # e.g., Cologne [LOC] vs Köln [LOC] ≈ 0.818; Apple [ORG] vs Microsoft [ORG] ≈ 0.400; Paris [LOC] vs Berlin [LOC] ≈ 0.491
135 )
136 for i, li in enumerate(all_labels):
137 for j, lj in enumerate(all_labels):
138 print(f"{li:20s} vs {lj:20s}: {sim[i, j]:.3f}")
139 print()
140
141# Within-document similarity matrices
142print("\nPer-document similarity (only if >=2 entities predicted):") # e.g., Doc 6: Cologne vs Cologne = 0.986
143for r in results:
144 embs = r["embeddings"]
145 if len(embs) < 2:
146 continue
147 sim_doc = embs @ embs.T
148 print(f"\nText: {r['text']}")
149 for i, label_i in enumerate(r["mention_surfaces"]):
150 for j, label_j in enumerate(r["mention_surfaces"]):
151 print(f" {label_i:12s} vs {label_j:12s}: {sim_doc[i, j]:.3f}")| Dataset | Split | Precision | Recall | F1 | Support |
|---|---|---|---|---|---|
| CoNLL 2003 (en) | test | 0.9111 | 0.9551 | 0.9326 | 4,946 |
| CoNLL 2002 (es) | test | 0.7515 | 0.8568 | 0.8007 | 3,219 |
| GermEval 2014 (de) | test | 0.8290 | 0.8663 | 0.8473 | 3,157 |
Note: The released checkpoint predicts the coarse types {PER, ORG, LOC}. Historical "MISC" mentions are mapped toOduring decoding; evaluation scripts should ignore that label when computing metrics.
| Language | Docs |
|---|---|
| en | 68,809 |
| zh | 12,926 |
| pt | 12,258 |
| sk | 11,361 |
| hr | 9,842 |
| sv | 8,133 |
| da | 6,537 |
| sr | 5,192 |
| de | 3,080 |
| fr | 2,584 |
| ru | 2,448 |
| es | 2,413 |
| it | 1,827 |
| ja | 1,579 |
| ko | 1,454 |
| nl | 1,317 |
| ar | 1,290 |
| pl | 1,228 |
| tl | 1,192 |
| cs | 1,161 |
| tr | 1,156 |
| no | 1,115 |
| uk | 1,100 |
| fi | 1,062 |
| vi | 1,018 |
| ro | 1,017 |
| id | 977 |
| lv | 950 |
| ms | 928 |
| el | 905 |
| bg | 901 |
| bn | 893 |
| fa | 876 |
| pa | 845 |
| ta | 845 |
| th | 834 |
| sl | 832 |
| hu | 829 |
| he | 822 |
| te | 791 |
| hi | 780 |
| mr | 774 |
| lt | 765 |
| ur | 745 |
| et | 734 |
| ml | 712 |
| gu | 629 |
| ca | 554 |
| jv | 534 |
| sw | 465 |
| my | 450 |
| az | 436 |
| ceb | 188 |