Upload folder using huggingface_hub

modeling_xlmedu.py

@@ -9,12 +9,11 @@ Loading from the Hub (recommended):
     config = AutoConfig.from_pretrained("xaver-krueckl/situation-entity-segmenter", trust_remote_code=True)
     model  = AutoModel.from_pretrained("xaver-krueckl/situation-entity-segmenter", trust_remote_code=True)
 
-Inference:
+Inference (one pre-split sentence):
     from transformers import AutoTokenizer
     tokenizer = AutoTokenizer.from_pretrained("FacebookAI/xlm-roberta-large")
-    inputs    = tokenizer("Hello world.", return_tensors="pt")
-    tag_ids   = model.predict(inputs["input_ids"], inputs["attention_mask"])
-    tags      = [model.config.id2label[i] for i in tag_ids[0]]
+    words     = ["The", "cat", "sat", "."]
+    pairs     = model.predict_text(words, tokenizer)  # [(word, tag), ...]
 """
 
 from __future__ import annotations
@@ -104,8 +103,57 @@ class XLMEduModelHF(PreTrainedModel):
         Returns one integer tag per sub-token (including special tokens).
         During training, non-first sub-tokens were masked (-100 → O), so their
         predictions are meaningless. Callers should use encoding.word_ids() to
-        read only the first sub-token of each word. See the model card for a
-        complete inference example.
+        read only the first sub-token of each word.
         """
         emissions = self.forward(input_ids, attention_mask)["logits"]
         return self.crf.decode(emissions.float(), mask=attention_mask.bool())
+
+    def predict_text(
+        self,
+        words:     List[str],
+        tokenizer,
+    ) -> List[tuple]:
+        """Tag a single pre-tokenised sentence.
+
+        Args:
+            words:     word-level tokens for one sentence (e.g. from spaCy).
+            tokenizer: the HuggingFace tokenizer for this model.
+
+        Returns:
+            List of (word, tag) pairs, one per input word.
+        """
+        from collections import defaultdict
+
+        encoding = tokenizer(
+            words,
+            return_tensors="pt",
+            is_split_into_words=True,
+            truncation=True,
+            max_length=512,
+        )
+
+        with torch.no_grad():
+            raw_tag_ids = self.predict(
+                encoding["input_ids"], encoding["attention_mask"]
+            )[0]
+
+        word_ids       = encoding.word_ids(batch_index=0)
+        word_token_ids: dict = defaultdict(list)
+        for pos, word_idx in enumerate(word_ids):
+            if word_idx is not None:
+                word_token_ids[word_idx].append(
+                    encoding["input_ids"][0, pos].item()
+                )
+
+        results: List[tuple] = []
+        prev_word_idx = None
+        for pos, word_idx in enumerate(word_ids):
+            if word_idx is None:
+                continue
+            if word_idx != prev_word_idx:
+                word = tokenizer.decode(word_token_ids[word_idx]).strip()
+                tag  = self.config.id2label[raw_tag_ids[pos]]
+                results.append((word, tag))
+                prev_word_idx = word_idx
+
+        return results