add pipeline_component

Author: hiroshi-matsuda-rit

config/ja_electra_base_parser_ner_accuracy.cfg

@@ -66,7 +66,7 @@ pooling = {"@layers":"reduce_mean.v1"}
 upstream = "*"
 
 [components.transformer]
-factory = "transformer"
+factory = "transformer_custom"
 max_batch_items = 4096
 set_extra_annotations = {"@annotation_setters":"spacy-transformers.null_annotation_setter.v1"}
 

ginza_transformers/pipeline_component.py

@@ -0,0 +1,78 @@
+from typing import List, Callable, Iterable, Union
+from pathlib import Path
+
+from spacy.language import Language
+from spacy.pipeline.pipe import deserialize_config
+from spacy.tokens import Doc
+from spacy import util
+from thinc.api import Model, Config
+
+from spacy_transformers.data_classes import FullTransformerBatch
+from spacy_transformers.pipeline_component import Transformer, DOC_EXT_ATTR
+
+from ginza_transformers.util import huggingface_from_pretrained_custom
+
+
+DEFAULT_CONFIG_STR = """
+[transformer_custom]
+max_batch_items = 4096
+
+[transformer_custom.set_extra_annotations]
+@annotation_setters = "spacy-transformers.null_annotation_setter.v1"
+
+[transformer_custom.model]
+@architectures = "ginza-transformers.TransformerModel.v1"
+name = "electra-base-ud-japanese-discriminator"
+tokenizer_config = {"use_fast": false, "tokenizer_class": "sudachitra.tokenization_electra_sudachipy.ElectraSudachipyTokenizer"}
+
+[transformer_custom.model.get_spans]
+@span_getters = "spacy-transformers.strided_spans.v1"
+window = 128
+stride = 96
+"""
+
+DEFAULT_CONFIG = Config().from_str(DEFAULT_CONFIG_STR)
+
+
+@Language.factory(
+    "transformer_custom",
+    assigns=[f"doc._.{DOC_EXT_ATTR}"],
+    default_config=DEFAULT_CONFIG["transformer_custom"],
+)
+def make_transformer_custom(
+    nlp: Language,
+    name: str,
+    model: Model[List[Doc], FullTransformerBatch],
+    set_extra_annotations: Callable[[List[Doc], FullTransformerBatch], None],
+    max_batch_items: int,
+):
+    return TransformerCustom(
+        nlp.vocab,
+        model,
+        set_extra_annotations,
+        max_batch_items=max_batch_items,
+        name=name,
+    )
+
+
+class TransformerCustom(Transformer):
+
+    def from_disk(
+        self, path: Union[str, Path], *, exclude: Iterable[str] = tuple()
+    ) -> "TransformerCustom":
+
+        def load_model(p):
+            p = Path(p).absolute()
+            tokenizer, transformer = huggingface_from_pretrained_custom(
+                p, self.model.attrs["tokenizer_config"]
+            )
+            self.model.attrs["tokenizer"] = tokenizer
+            self.model.attrs["set_transformer"](self.model, transformer)
+
+        deserialize = {
+            "vocab": self.vocab.from_disk,
+            "cfg": lambda p: self.cfg.update(deserialize_config(p)),
+            "model": load_model,
+        }
+        util.from_disk(path, deserialize, exclude)
+        return self

setup.py

@@ -6,6 +6,9 @@
     author_email="ginza@megagon.ai",
     description="ginza-transformers",
     entry_points={
+        "spacy_factories": [
+            "transformer_custom = ginza_transformers.pipeline_component:make_transformer_custom",
+        ],
         "spacy_architectures": [
             "ginza-transformers.TransformerModel.v1 = ginza_transformers:architectures.TransformerModelCustom",
         ],
@@ -17,5 +20,5 @@
     name="ginza-transformers",
     packages=find_packages(include=["ginza_transformers", "ginza_transformers.layers"]),
     url="https://github.com/megagonlabs/ginza-transformers",
-    version='0.1.1',
+    version='0.2.0',
 )