[ML] Improve NLP model import by using nicely defined types (#459)

This adds some more definite types for our NLP tasks and tokenization configurations.

This is the first step in allowing users to more easily import their own transformer models via something other than hugging face.

Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>

README.md

@@ -250,11 +250,11 @@ Downloading: 100%|██████████| 249M/249M [00:23<00:00, 11.2MB
 # Export the model in a TorchScrpt representation which Elasticsearch uses
 >>> tmp_path = "models"
 >>> Path(tmp_path).mkdir(parents=True, exist_ok=True)
->>> model_path, config_path, vocab_path = tm.save(tmp_path)
+>>> model_path, config, vocab_path = tm.save(tmp_path)
 
 # Import model into Elasticsearch
 >>> es = elasticsearch.Elasticsearch("http://elastic:mlqa_admin@localhost:9200", timeout=300)  # 5 minute timeout
 >>> ptm = PyTorchModel(es, tm.elasticsearch_model_id())
->>> ptm.import_model(model_path, config_path, vocab_path)
+>>> ptm.import_model(model_path=model_path, config_path=None, vocab_path=vocab_path, config=config)
 100%|██████████| 63/63 [00:12<00:00,  5.02it/s]
 ```

bin/eland_import_hub_model

@@ -188,7 +188,7 @@ if __name__ == "__main__":
         logger.info(f"Loading HuggingFace transformer tokenizer and model '{args.hub_model_id}'")
 
         tm = TransformerModel(args.hub_model_id, args.task_type, args.quantize)
-        model_path, config_path, vocab_path = tm.save(tmp_dir)
+        model_path, config, vocab_path = tm.save(tmp_dir)
 
         ptm = PyTorchModel(es, args.es_model_id if args.es_model_id else tm.elasticsearch_model_id())
         model_exists = es.options(ignore_status=404).ml.get_trained_models(model_id=ptm.model_id).meta.status == 200
@@ -206,7 +206,7 @@ if __name__ == "__main__":
                 exit(1)
 
         logger.info(f"Creating model with id '{ptm.model_id}'")
-        ptm.put_config(config_path)
+        ptm.put_config(config=config)
 
         logger.info(f"Uploading model definition")
         ptm.put_model(model_path)

eland/ml/pytorch/_pytorch_model.py

@@ -19,11 +19,22 @@
 import json
 import math
 import os
-from typing import TYPE_CHECKING, Any, Iterable, List, Mapping, Set, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
 
 from tqdm.auto import tqdm  # type: ignore
 
 from eland.common import ensure_es_client
+from eland.ml.pytorch.nlp_ml_model import NlpTrainedModelConfig
 
 if TYPE_CHECKING:
     from elasticsearch import Elasticsearch
@@ -49,10 +60,19 @@ def __init__(
         self._client: Elasticsearch = ensure_es_client(es_client)
         self.model_id = model_id
 
-    def put_config(self, path: str) -> None:
-        with open(path) as f:
-            config = json.load(f)
-        self._client.ml.put_trained_model(model_id=self.model_id, **config)
+    def put_config(
+        self, path: Optional[str] = None, config: Optional[NlpTrainedModelConfig] = None
+    ) -> None:
+        if path is not None and config is not None:
+            raise ValueError("Only include path or config. Not both")
+        if path is not None:
+            with open(path) as f:
+                config_map = json.load(f)
+        elif config is not None:
+            config_map = config.to_dict()
+        else:
+            raise ValueError("Must provide path or config")
+        self._client.ml.put_trained_model(model_id=self.model_id, **config_map)
 
     def put_vocab(self, path: str) -> None:
         with open(path) as f:
@@ -89,13 +109,14 @@ def model_file_chunk_generator() -> Iterable[str]:
 
     def import_model(
         self,
+        *,
         model_path: str,
-        config_path: str,
+        config_path: Optional[str],
         vocab_path: str,
+        config: Optional[NlpTrainedModelConfig] = None,
         chunk_size: int = DEFAULT_CHUNK_SIZE,
     ) -> None:
-        # TODO: Implement some pre-flight checks on config, vocab, and model
-        self.put_config(config_path)
+        self.put_config(path=config_path, config=config)
         self.put_model(model_path, chunk_size)
         self.put_vocab(vocab_path)
 

eland/ml/pytorch/nlp_ml_model.py

@@ -0,0 +1,228 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+import typing as t
+
+
+class NlpTokenizationConfig:
+    def __init__(self, *, configuration_type: str):
+        self.name = configuration_type
+
+    def to_dict(self):
+        return {
+            self.name: {
+                k: v for k, v in self.__dict__.items() if v is not None and k != "name"
+            }
+        }
+
+
+class NlpRobertaTokenizationConfig(NlpTokenizationConfig):
+    def __init__(
+        self,
+        *,
+        add_prefix_space: t.Optional[bool] = None,
+        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="roberta")
+        self.add_prefix_space = add_prefix_space
+        self.with_special_tokens = with_special_tokens
+        self.max_sequence_length = max_sequence_length
+        self.truncate = truncate
+        self.span = span
+
+
+class NlpBertTokenizationConfig(NlpTokenizationConfig):
+    def __init__(
+        self,
+        *,
+        do_lower_case: t.Optional[bool] = None,
+        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="bert")
+        self.do_lower_case = do_lower_case
+        self.with_special_tokens = with_special_tokens
+        self.max_sequence_length = max_sequence_length
+        self.truncate = truncate
+        self.span = span
+
+
+class NlpMPNetTokenizationConfig(NlpTokenizationConfig):
+    def __init__(
+        self,
+        *,
+        do_lower_case: t.Optional[bool] = None,
+        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="mpnet")
+        self.do_lower_case = do_lower_case
+        self.with_special_tokens = with_special_tokens
+        self.max_sequence_length = max_sequence_length
+        self.truncate = truncate
+        self.span = span
+
+
+class InferenceConfig:
+    def __init__(self, *, configuration_type: str):
+        self.name = configuration_type
+
+    def to_dict(self) -> t.Dict[str, t.Any]:
+        return {
+            self.name: {
+                k: v.to_dict() if hasattr(v, "to_dict") else v
+                for k, v in self.__dict__.items()
+                if v is not None and k != "name"
+            }
+        }
+
+
+class TextClassificationInferenceOptions(InferenceConfig):
+    def __init__(
+        self,
+        *,
+        classification_labels: t.Union[t.List[str], t.Tuple[str, ...]],
+        tokenization: NlpTokenizationConfig,
+        results_field: t.Optional[str] = None,
+        num_top_classes: t.Optional[int] = None,
+    ):
+        super().__init__(configuration_type="text_classification")
+        self.results_field = results_field
+        self.num_top_classes = num_top_classes
+        self.tokenization = tokenization
+        self.classification_labels = classification_labels
+
+
+class ZeroShotClassificationInferenceOptions(InferenceConfig):
+    def __init__(
+        self,
+        *,
+        tokenization: NlpTokenizationConfig,
+        classification_labels: t.Union[t.List[str], t.Tuple[str, ...]],
+        results_field: t.Optional[str] = None,
+        multi_label: t.Optional[bool] = None,
+        labels: t.Optional[t.Union[t.List[str], t.Tuple[str, ...]]] = None,
+        hypothesis_template: t.Optional[str] = None,
+    ):
+        super().__init__(configuration_type="zero_shot_classification")
+        self.tokenization = tokenization
+        self.hypothesis_template = hypothesis_template
+        self.classification_labels = classification_labels
+        self.results_field = results_field
+        self.multi_label = multi_label
+        self.labels = labels
+
+
+class FillMaskInferenceOptions(InferenceConfig):
+    def __init__(
+        self,
+        *,
+        tokenization: NlpTokenizationConfig,
+        results_field: t.Optional[str] = None,
+        num_top_classes: t.Optional[int] = None,
+    ):
+        super().__init__(configuration_type="fill_mask")
+        self.num_top_classes = num_top_classes
+        self.tokenization = tokenization
+        self.results_field = results_field
+
+
+class NerInferenceOptions(InferenceConfig):
+    def __init__(
+        self,
+        *,
+        tokenization: NlpTokenizationConfig,
+        classification_labels: t.Union[t.List[str], t.Tuple[str, ...]],
+        results_field: t.Optional[str] = None,
+    ):
+        super().__init__(configuration_type="ner")
+        self.tokenization = tokenization
+        self.classification_labels = classification_labels
+        self.results_field = results_field
+
+
+class PassThroughInferenceOptions(InferenceConfig):
+    def __init__(
+        self,
+        *,
+        tokenization: NlpTokenizationConfig,
+        results_field: t.Optional[str] = None,
+    ):
+        super().__init__(configuration_type="pass_through")
+        self.tokenization = tokenization
+        self.results_field = results_field
+
+
+class TextEmbeddingInferenceOptions(InferenceConfig):
+    def __init__(
+        self,
+        *,
+        tokenization: NlpTokenizationConfig,
+        results_field: t.Optional[str] = None,
+    ):
+        super().__init__(configuration_type="text_embedding")
+        self.tokenization = tokenization
+        self.results_field = results_field
+
+
+class TrainedModelInput:
+    def __init__(self, *, field_names: t.List[str]):
+        self.field_names = field_names
+
+    def to_dict(self) -> t.Dict[str, t.Any]:
+        return self.__dict__
+
+
+class NlpTrainedModelConfig:
+    def __init__(
+        self,
+        *,
+        description: str,
+        inference_config: InferenceConfig,
+        input: TrainedModelInput = TrainedModelInput(field_names=["text_field"]),
+        metadata: t.Optional[dict] = None,
+        model_type: t.Union["t.Literal['pytorch']", str] = "pytorch",
+        default_field_map: t.Optional[t.Mapping[str, str]] = None,
+        tags: t.Optional[t.Union[t.List[str], t.Tuple[str, ...]]] = None,
+    ):
+        self.tags = tags
+        self.default_field_map = default_field_map
+        self.description = description
+        self.inference_config = inference_config
+        self.input = input
+        self.metadata = metadata
+        self.model_type = model_type
+
+    def to_dict(self) -> t.Dict[str, t.Any]:
+        return {
+            k: v.to_dict() if hasattr(v, "to_dict") else v
+            for k, v in self.__dict__.items()
+            if v is not None
+        }