change strategy: overriding HFShim.from_bytes

Author: hiroshi-matsuda-rit

ginza_transformers/layers/hf_shim_custom.py

@@ -0,0 +1,101 @@
+import sys
+from typing import Any
+from io import BytesIO
+from pathlib import Path
+import srsly
+import torch
+from spacy.util import SimpleFrozenDict
+from spacy.vectors import get_current_ops
+
+from spacy_transformers.layers import hf_shim
+from spacy_transformers.layers.hf_shim import HFShim
+from spacy_transformers.data_classes import HFObjects
+from spacy_transformers.util import make_tempdir
+
+from transformers import AutoModel, AutoConfig, AutoTokenizer
+
+
+def override_hf_shims_from_bytes():
+    assert hf_shim.HFShim.from_bytes is not HFShimCustom.from_bytes
+    origin = hf_shim.HFShim.from_bytes
+    hf_shim.HFShim.from_bytes = HFShimCustom.from_bytes
+    return origin
+
+def recover_hf_shims_from_bytes(origin):
+    assert hf_shim.HFShim.from_bytes is HFShimCustom.from_bytes
+    hf_shim.HFShim.from_bytes = origin
+
+
+class HFShimCustom(HFShim):
+
+    def from_bytes(self, bytes_data):
+        msg = srsly.msgpack_loads(bytes_data)
+        config_dict = msg["config"]
+        tok_dict = msg["tokenizer"]
+        if config_dict:
+            with make_tempdir() as temp_dir:
+                config_file = temp_dir / "config.json"
+                srsly.write_json(config_file, config_dict)
+                config = AutoConfig.from_pretrained(config_file)
+                for x, x_bytes in tok_dict.items():
+                    Path(temp_dir / x).write_bytes(x_bytes)
+                tokenizer = None
+                try:
+                    tokenizer = AutoTokenizer.from_pretrained(str(temp_dir.absolute()))
+                except (ValueError, OSError):
+                    pass
+                if tokenizer is None:
+                    tok_config = srsly.read_json(str((temp_dir / "tokenizer_config.json").absolute()))
+                    tokenizer_class_name = tok_config["tokenizer_class"].split(".")
+                    if tokenizer_class_name == ["ElectraSudachipyTokenizer"]:
+                        from sudachitra.tokenization_electra_sudachipy import ElectraSudachipyTokenizer as tokenizer_class
+                        tokenizer = tokenizer_class(vocab_file=str((temp_dir / "vocab.txt").absolute()), **tok_config)
+                    else:
+                        from importlib import import_module
+                        tokenizer_module = import_module(".".join(tokenizer_class_name[:-1]))
+                        tokenizer_class = getattr(tokenizer_module, tokenizer_class_name[-1])
+
+                vocab_file_contents = None
+                if hasattr(tokenizer, "vocab_file"):
+                    vocab_file_name = tokenizer.vocab_files_names["vocab_file"]
+                    vocab_file_path = str((temp_dir / vocab_file_name).absolute())
+                    with open(vocab_file_path, "rb") as fileh:
+                        vocab_file_contents = fileh.read()
+
+            try:
+                transformer = AutoModel.from_config(config)
+            except OSError as e:
+                try:
+                    transformer = AutoModel.from_pretrained(config["_name_or_path"], local_files_only=True)
+                except OSError as e2:
+                    print("trying to download model from huggingface hub:", config["_name_or_path"], "...", file=sys.stderr)
+                    transformer = AutoModel.from_pretrained(config["_name_or_path"])
+                    print("succeded", file=sys.stderr)
+
+            self._hfmodel = HFObjects(
+                tokenizer,
+                transformer,
+                vocab_file_contents,
+                SimpleFrozenDict(),
+                SimpleFrozenDict(),
+            )
+            self._model = transformer
+            filelike = BytesIO(msg["state"])
+            filelike.seek(0)
+            ops = get_current_ops()
+            if ops.device_type == "cpu":
+                map_location = "cpu"
+            else:  # pragma: no cover
+                device_id = torch.cuda.current_device()
+                map_location = f"cuda:{device_id}"
+            self._model.load_state_dict(torch.load(filelike, map_location=map_location))
+            self._model.to(map_location)
+        else:
+            self._hfmodel = HFObjects(
+                None,
+                None,
+                None,
+                msg["_init_tokenizer_config"],
+                msg["_init_transformer_config"],
+            )
+        return self

ginza_transformers/layers/transformer_model.py

@@ -1,184 +1,5 @@
-import copy
-import sys
-from typing import Callable, Dict, Optional, Tuple, Union
-from pathlib import Path
+from spacy_transformers.layers.transformer_model import TransformerModel
 
-from transformers import AutoConfig, AutoModel, AutoTokenizer, PreTrainedTokenizerBase
 
-from thinc.api import CupyOps, Model, get_current_ops
-
-from spacy_transformers.align import get_alignment
-from spacy_transformers.data_classes import WordpieceBatch, HFObjects
-from spacy_transformers.layers._util import replace_listener, replace_listener_cfg
-from spacy_transformers.layers.hf_wrapper import HFWrapper
-from spacy_transformers.layers.transformer_model import (
-    TransformerModel,
-    _convert_transformer_inputs,
-    _convert_transformer_outputs,
-    forward,
-    huggingface_tokenize,
-    set_pytorch_transformer,
-)
-from spacy_transformers.truncate import truncate_oversize_splits
-
-
-class TransformerModelCustom(Model):
-    def __init__(
-        self,
-        name: str,
-        get_spans: Callable,
-        tokenizer_config: dict = {},
-        transformer_config: dict = {},
-        mixed_precision: bool = False,
-        grad_scaler_config: dict = {},
-    ):
-        """
-        get_spans (Callable[[List[Doc]], List[Span]]):
-            A function to extract spans from the batch of Doc objects.
-            This is used to manage long documents, by cutting them into smaller
-            sequences before running the transformer. The spans are allowed to
-            overlap, and you can also omit sections of the Doc if they are not
-            relevant.
-        tokenizer_config (dict): Settings to pass to the transformers tokenizer.
-        transformer_config (dict): Settings to pass to the transformers forward pass.
-        """
-        hf_model = HFObjects(None, None, None, tokenizer_config, transformer_config)
-        wrapper = HFWrapper(
-            hf_model,
-            convert_inputs=_convert_transformer_inputs,
-            convert_outputs=_convert_transformer_outputs,
-            mixed_precision=mixed_precision,
-            grad_scaler_config=grad_scaler_config,
-        )
-        super().__init__(
-            "transformer",
-            forward,
-            init=init_custom,
-            layers=[wrapper],
-            dims={"nO": None},
-            attrs={
-                "get_spans": get_spans,
-                "name": name,
-                "set_transformer": set_pytorch_transformer,
-                "has_transformer": False,
-                "flush_cache_chance": 0.0,
-                "replace_listener": replace_listener,
-                "replace_listener_cfg": replace_listener_cfg,
-            },
-        )
-
-    @property
-    def tokenizer(self):
-        return self.layers[0].shims[0]._hfmodel.tokenizer
-
-    @property
-    def transformer(self):
-        return self.layers[0].shims[0]._hfmodel.transformer
-
-    @property
-    def _init_tokenizer_config(self):
-        return self.layers[0].shims[0]._hfmodel._init_tokenizer_config
-
-    @property
-    def _init_transformer_config(self):
-        return self.layers[0].shims[0]._hfmodel._init_transformer_config
-
-    def copy(self):
-        """
-        Create a copy of the model, its attributes, and its parameters. Any child
-        layers will also be deep-copied. The copy will receive a distinct `model.id`
-        value.
-        """
-        copied = TransformerModel(self.name, self.attrs["get_spans"])
-        params = {}
-        for name in self.param_names:
-            params[name] = self.get_param(name) if self.has_param(name) else None
-        copied.params = copy.deepcopy(params)
-        copied.dims = copy.deepcopy(self._dims)
-        copied.layers[0] = copy.deepcopy(self.layers[0])
-        for name in self.grad_names:
-            copied.set_grad(name, self.get_grad(name).copy())
-        return copied
-
-
-def init_custom(model: Model, X=None, Y=None):
-    if model.attrs["has_transformer"]:
-        return
-    name = model.attrs["name"]
-    tok_cfg = model._init_tokenizer_config
-    trf_cfg = model._init_transformer_config
-    tokenizer, hf_model = huggingface_from_pretrained_custom(name, tok_cfg, trf_cfg, model.attrs["name"])
-    model.attrs["set_transformer"](model, hf_model)
-    # Call the model with a batch of inputs to infer the width
-    if X:
-        # If we're dealing with actual texts, do the work to setup the wordpieces
-        # batch properly
-        docs = X
-        get_spans = model.attrs["get_spans"]
-        nested_spans = get_spans(docs)
-        flat_spans = []
-        for doc_spans in nested_spans:
-            flat_spans.extend(doc_spans)
-        token_data = huggingface_tokenize(tokenizer, [span.text for span in flat_spans])
-        wordpieces = WordpieceBatch.from_batch_encoding(token_data)
-        align = get_alignment(
-            flat_spans, wordpieces.strings, tokenizer.all_special_tokens
-        )
-        wordpieces, align = truncate_oversize_splits(
-            wordpieces, align, tokenizer.model_max_length
-        )
-    else:
-        texts = ["hello world", "foo bar"]
-        token_data = huggingface_tokenize(tokenizer, texts)
-        wordpieces = WordpieceBatch.from_batch_encoding(token_data)
-    model.layers[0].initialize(X=wordpieces)
-    model_output = model.layers[0].predict(wordpieces)
-    model.set_dim("nO", model_output.last_hidden_state.shape[-1])
-
-
-def huggingface_from_pretrained_custom(
-    source: Union[Path, str], tok_config: Dict, trf_config: Dict, model_name: Optional[str] = None,
-) -> Tuple[PreTrainedTokenizerBase, HFObjects]:
-    """Create a Huggingface transformer model from pretrained weights. Will
-    download the model if it is not already downloaded.
-
-    source (Union[str, Path]): The name of the model or a path to it, such as
-        'bert-base-cased'.
-    tok_config (dict): Settings to pass to the tokenizer.
-    trf_config (dict): Settings to pass to the transformer.
-    """
-    if hasattr(source, "absolute"):
-        str_path = str(source.absolute())
-    else:
-        str_path = source
-
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(str_path, **tok_config)
-    except ValueError as e:
-        if "tokenizer_class" not in tok_config:
-            raise e
-        tokenizer_class_name = tok_config["tokenizer_class"].split(".")
-        from importlib import import_module
-        tokenizer_module = import_module(".".join(tokenizer_class_name[:-1]))
-        tokenizer_class = getattr(tokenizer_module, tokenizer_class_name[-1])
-        tokenizer = tokenizer_class(vocab_file=str_path + "/vocab.txt", **tok_config)
-    vocab_file_contents = None
-    if hasattr(tokenizer, "vocab_file"):
-        with open(tokenizer.vocab_file, "rb") as fileh:
-            vocab_file_contents = fileh.read()
-
-    try:
-        trf_config["return_dict"] = True
-        config = AutoConfig.from_pretrained(str_path, **trf_config)
-        transformer = AutoModel.from_pretrained(str_path, config=config)
-    except OSError as e:
-        try:
-            transformer = AutoModel.from_pretrained(model_name, local_files_only=True)
-        except OSError as e2:
-            print("trying to download model from huggingface hub:", model_name, "...", file=sys.stderr)
-            transformer = AutoModel.from_pretrained(model_name)
-            print("succeded", file=sys.stderr)
-    ops = get_current_ops()
-    if isinstance(ops, CupyOps):
-        transformer.cuda()
-    return tokenizer, HFObjects(tokenizer, transformer, vocab_file_contents)
+class TransformerModelCustom(TransformerModel):
+    pass

ginza_transformers/pipeline_component.py

@@ -3,15 +3,13 @@
 from typing import List, Callable, Iterable, Union
 
 from spacy.language import Language
-from spacy.pipeline.pipe import deserialize_config
 from spacy.tokens import Doc
-from spacy import util, Errors
 from thinc.api import Model, Config
 
 from spacy_transformers.data_classes import FullTransformerBatch
 from spacy_transformers.pipeline_component import Transformer, DOC_EXT_ATTR
 
-from .layers.transformer_model import huggingface_from_pretrained_custom
+from .layers.hf_shim_custom import override_hf_shims_from_bytes, recover_hf_shims_from_bytes
 
 
 DEFAULT_CONFIG_STR = """
@@ -22,7 +20,7 @@
 @annotation_setters = "spacy-transformers.null_annotation_setter.v1"
 
 [transformer_custom.model]
-@architectures = "ginza-transformers.TransformerModel.v1"
+@architectures = "ginza-transformers.TransformerModel.v3"
 
 [transformer_custom.model.get_spans]
 @span_getters = "spacy-transformers.strided_spans.v1"
@@ -60,27 +58,9 @@ def from_disk(
         self, path: Union[str, Path], *, exclude: Iterable[str] = tuple()
     ) -> "TransformerCustom":
 
-        def load_model(p):
-            try:
-                with open(p, "rb") as mfile:
-                    self.model.from_bytes(mfile.read())
-            except AttributeError:
-                raise ValueError(Errors.E149) from None
-            except (IsADirectoryError, PermissionError):
-                p = Path(p).absolute()
-                tokenizer, hf_model = huggingface_from_pretrained_custom(
-                    p,
-                    self.model._init_tokenizer_config,
-                    self.model._init_transformer_config,
-                    self.model.attrs["name"],
-                )
-                self.model.attrs["tokenizer"] = tokenizer
-                self.model.attrs["set_transformer"](self.model, hf_model)
-
-        deserialize = {
-            "vocab": self.vocab.from_disk,
-            "cfg": lambda p: self.cfg.update(deserialize_config(p)),
-            "model": load_model,
-        }
-        util.from_disk(path, deserialize, exclude)
+        origin = override_hf_shims_from_bytes()
+        try:
+            super().from_disk(path, exclude=exclude)
+        finally:
+            recover_hf_shims_from_bytes(origin)
         return self