change strategy: overriding huggingface_from_pretrained

Author: hiroshi-matsuda-rit

ginza_transformers/layers/transformer_model.py

@@ -1,5 +1,70 @@
-from spacy_transformers.layers.transformer_model import TransformerModel
+from pathlib import Path
+import sys
+from typing import Union, Dict
 
+from transformers import AutoConfig, AutoModel, AutoTokenizer
 
-class TransformerModelCustom(TransformerModel):
-    pass
+from spacy_transformers.layers import transformer_model
+from spacy_transformers.data_classes import HFObjects
+
+from thinc.api import get_current_ops, CupyOps
+
+
+def override_huggingface_from_pretrained():
+    assert transformer_model.huggingface_from_pretrained is not huggingface_from_pretrained_custom
+    origin = transformer_model.huggingface_from_pretrained
+    transformer_model.huggingface_from_pretrained = huggingface_from_pretrained_custom
+    return origin
+
+def recover_huggingface_from_pretrained(origin):
+    assert transformer_model.huggingface_from_pretrained is huggingface_from_pretrained_custom
+    transformer_model.huggingface_from_pretrained = origin
+
+
+def huggingface_from_pretrained_custom(
+    source: Union[Path, str], tok_config: Dict, trf_config: Dict
+) -> 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(str_path, local_files_only=True)
+        except OSError as e2:
+            model_name = str(source)
+            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 HFObjects(tokenizer, transformer, vocab_file_contents)

ginza_transformers/pipeline_component.py

@@ -1,15 +1,17 @@
 from pathlib import Path
 import sys
-from typing import List, Callable, Iterable, Union
+from typing import List, Callable, Iterable, Optional, Union
 
 from spacy.language import Language
 from spacy.tokens import Doc
 from thinc.api import Model, Config
 
 from spacy_transformers.data_classes import FullTransformerBatch
 from spacy_transformers.pipeline_component import Transformer, DOC_EXT_ATTR
+from spacy.training import Example
 
 from .layers.hf_shim_custom import override_hf_shims_from_bytes, recover_hf_shims_from_bytes
+from .layers.transformer_model import override_huggingface_from_pretrained, recover_huggingface_from_pretrained
 
 
 DEFAULT_CONFIG_STR = """
@@ -54,10 +56,21 @@ def make_transformer_custom(
 
 class TransformerCustom(Transformer):
 
+    def initialize(
+        self,
+        get_examples: Callable[[], Iterable[Example]],
+        *,
+        nlp: Optional[Language] = None,
+    ):
+        origin = override_huggingface_from_pretrained()
+        try:
+            super().initialize(get_examples, nlp=nlp)
+        finally:
+            recover_huggingface_from_pretrained(origin)
+
     def from_disk(
         self, path: Union[str, Path], *, exclude: Iterable[str] = tuple()
     ) -> "TransformerCustom":
-
         origin = override_hf_shims_from_bytes()
         try:
             super().from_disk(path, exclude=exclude)