[ML] add ability to upload xlm-roberta tokenized models

eland/ml/pytorch/__init__.py

@@ -31,6 +31,7 @@
     "NlpTrainedModelConfig",
     "NlpBertTokenizationConfig",
     "NlpRobertaTokenizationConfig",
+    "NlpXLMRobertaTokenizationConfig",
     "NlpMPNetTokenizationConfig",
     "task_type_from_model_config",
 ]

eland/ml/pytorch/nlp_ml_model.py

@@ -66,6 +66,26 @@ def __init__(
         self.add_prefix_space = add_prefix_space
 
 
+class NlpXLMRobertaTokenizationConfig(NlpTokenizationConfig):
+    def __init__(
+        self,
+        *,
+        with_special_tokens: t.Optional[bool] = None,
+        max_sequence_length: t.Optional[int] = None,
+        truncate: t.Optional[
+            t.Union["t.Literal['first', 'none', 'second']", str]
+        ] = None,
+        span: t.Optional[int] = None,
+    ):
+        super().__init__(
+            configuration_type="xlm_roberta",
+            with_special_tokens=with_special_tokens,
+            max_sequence_length=max_sequence_length,
+            truncate=truncate,
+            span=span,
+        )
+
+
 class NlpBertTokenizationConfig(NlpTokenizationConfig):
     def __init__(
         self,

eland/ml/pytorch/transformers.py

@@ -48,6 +48,7 @@
     NlpRobertaTokenizationConfig,
     NlpTokenizationConfig,
     NlpTrainedModelConfig,
+    NlpXLMRobertaTokenizationConfig,
     PassThroughInferenceOptions,
     QuestionAnsweringInferenceOptions,
     TextClassificationInferenceOptions,
@@ -108,6 +109,7 @@
     transformers.RobertaTokenizer,
     transformers.BartTokenizer,
     transformers.SqueezeBertTokenizer,
+    transformers.XLMRobertaTokenizer,
 )
 SUPPORTED_TOKENIZERS_NAMES = ", ".join(sorted([str(x) for x in SUPPORTED_TOKENIZERS]))
 
@@ -183,6 +185,7 @@ def from_pretrained(model_id: str) -> Optional[Any]:
             model.config,
             (
                 transformers.MPNetConfig,
+                transformers.XLMRobertaConfig,
                 transformers.RobertaConfig,
                 transformers.BartConfig,
             ),
@@ -295,6 +298,7 @@ def from_pretrained(
                 model.config,
                 (
                     transformers.MPNetConfig,
+                    transformers.XLMRobertaConfig,
                     transformers.RobertaConfig,
                     transformers.BartConfig,
                 ),
@@ -463,6 +467,7 @@ def _compatible_inputs(self) -> Tuple[Tensor, ...]:
             self._model.config,
             (
                 transformers.MPNetConfig,
+                transformers.XLMRobertaConfig,
                 transformers.RobertaConfig,
                 transformers.BartConfig,
             ),
@@ -639,6 +644,20 @@ def _load_vocab(self) -> Dict[str, List[str]]:
                 " ".join(m) for m, _ in sorted(ranks.items(), key=lambda kv: kv[1])
             ]
             vocab_obj["merges"] = merges
+        sp_model = getattr(self._tokenizer, "sp_model", None)
+        if sp_model:
+            id_correction = getattr(self._tokenizer, "fairseq_offset", 0)
+            scores = []
+            for _ in range(0, id_correction):
+                scores.append(0.0)
+            for token_id in range(id_correction, len(vocabulary)):
+                try:
+                    scores.append(sp_model.get_score(token_id - id_correction))
+                except IndexError:
+                    scores.append(0.0)
+                    pass
+            if len(scores) > 0:
+                vocab_obj["scores"] = scores
         return vocab_obj
 
     def _create_tokenization_config(self) -> NlpTokenizationConfig:
@@ -658,6 +677,12 @@ def _create_tokenization_config(self) -> NlpTokenizationConfig:
                     self._tokenizer, "max_model_input_sizes", dict()
                 ).get(self._model_id),
             )
+        elif isinstance(self._tokenizer, transformers.XLMRobertaTokenizer):
+            return NlpXLMRobertaTokenizationConfig(
+                max_sequence_length=getattr(
+                    self._tokenizer, "max_model_input_sizes", dict()
+                ).get(self._model_id),
+            )
         else:
             return NlpBertTokenizationConfig(
                 do_lower_case=getattr(self._tokenizer, "do_lower_case", None),

requirements-dev.txt

@@ -14,7 +14,7 @@ tqdm<5
 #
 scikit-learn>=0.22.1,<2
 xgboost>=0.90,<2
-lightgbm>=2,<4
+#lightgbm>=2,<4
 
 # PyTorch doesn't support Python 3.11 yet (pytorch/pytorch#86566)