Update pytorch-partial-tagger (#35)

* Remove unnecessary code

* Update dependencies

* Update incompatible code

* Bump version

* Update requirements.txt

Author: yasufumy

pyproject.toml

@@ -8,7 +8,7 @@ requires-python = ">=3.8"
 
 [tool.poetry]
 name = "spacy-partial-tagger"
-version = "0.14.0"
+version = "0.15.0"
 description = "Sequence Tagger for Partially Annotated Dataset in spaCy"
 authors = ["yasufumi <yasufumi.taniguchi@gmail.com>"]
 license = "MIT"
@@ -27,7 +27,7 @@ transformers = {extras = ["ja"], version = "^4.25.1"}
 torch = "^2.0.1"
 spacy = {extras = ["transformers"], version = "^3.3.1"}
 spacy-alignments = "^0.8.5"
-pytorch-partial-tagger = "^0.1.6"
+pytorch-partial-tagger = "^0.1.7"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.3.0"

requirements.txt

@@ -9,9 +9,9 @@ confection==0.0.4 ; python_version >= "3.8" and python_version < "4.0"
 coverage[toml]==7.2.7 ; python_version >= "3.8" and python_version < "4.0"
 cymem==2.0.7 ; python_version >= "3.8" and python_version < "4.0"
 exceptiongroup==1.1.1 ; python_version >= "3.8" and python_version < "3.11"
-filelock==3.12.0 ; python_version >= "3.8" and python_version < "4.0"
+filelock==3.12.1 ; python_version >= "3.8" and python_version < "4.0"
 flake8==4.0.1 ; python_version >= "3.8" and python_version < "4.0"
-fsspec==2023.5.0 ; python_version >= "3.8" and python_version < "4.0"
+fsspec==2023.6.0 ; python_version >= "3.8" and python_version < "4.0"
 fugashi==1.2.1 ; python_version >= "3.8" and python_version < "4.0"
 huggingface-hub==0.15.1 ; python_version >= "3.8" and python_version < "4.0"
 idna==3.4 ; python_version >= "3.8" and python_version < "4.0"
@@ -32,15 +32,15 @@ packaging==23.1 ; python_version >= "3.8" and python_version < "4.0"
 pathspec==0.11.1 ; python_version >= "3.8" and python_version < "4.0"
 pathy==0.10.1 ; python_version >= "3.8" and python_version < "4.0"
 plac==1.3.5 ; python_version >= "3.8" and python_version < "4.0"
-platformdirs==3.5.1 ; python_version >= "3.8" and python_version < "4.0"
+platformdirs==3.5.3 ; python_version >= "3.8" and python_version < "4.0"
 pluggy==1.0.0 ; python_version >= "3.8" and python_version < "4.0"
 preshed==3.0.8 ; python_version >= "3.8" and python_version < "4.0"
 pycodestyle==2.8.0 ; python_version >= "3.8" and python_version < "4.0"
-pydantic==1.10.8 ; python_version >= "3.8" and python_version < "4.0"
+pydantic==1.10.9 ; python_version >= "3.8" and python_version < "4.0"
 pyflakes==2.4.0 ; python_version >= "3.8" and python_version < "4.0"
 pytest-cov==3.0.0 ; python_version >= "3.8" and python_version < "4.0"
-pytest==7.3.1 ; python_version >= "3.8" and python_version < "4.0"
-pytorch-partial-tagger==0.1.6 ; python_version >= "3.8" and python_version < "4.0"
+pytest==7.3.2 ; python_version >= "3.8" and python_version < "4.0"
+pytorch-partial-tagger==0.1.7 ; python_version >= "3.8" and python_version < "4.0"
 pyyaml==6.0 ; python_version >= "3.8" and python_version < "4.0"
 regex==2023.6.3 ; python_version >= "3.8" and python_version < "4.0"
 requests==2.31.0 ; python_version >= "3.8" and python_version < "4.0"
@@ -51,6 +51,7 @@ smart-open==6.3.0 ; python_version >= "3.8" and python_version < "4.0"
 spacy-alignments==0.8.6 ; python_version >= "3.8" and python_version < "4.0"
 spacy-legacy==3.0.12 ; python_version >= "3.8" and python_version < "4.0"
 spacy-loggers==1.0.4 ; python_version >= "3.8" and python_version < "4.0"
+spacy-transformers==1.2.4 ; python_version >= "3.8" and python_version < "4.0"
 spacy==3.5.3 ; python_version >= "3.8" and python_version < "4.0"
 spacy[transformers]==3.5.3 ; python_version >= "3.8" and python_version < "4.0"
 srsly==2.4.6 ; python_version >= "3.8" and python_version < "4.0"
@@ -68,5 +69,5 @@ typer==0.7.0 ; python_version >= "3.8" and python_version < "4.0"
 typing-extensions==4.6.3 ; python_version >= "3.8" and python_version < "4.0"
 unidic-lite==1.0.8 ; python_version >= "3.8" and python_version < "4.0"
 unidic==1.1.0 ; python_version >= "3.8" and python_version < "4.0"
-urllib3==2.0.2 ; python_version >= "3.8" and python_version < "4.0"
+urllib3==2.0.3 ; python_version >= "3.8" and python_version < "4.0"
 wasabi==0.10.1 ; python_version >= "3.8" and python_version < "4.0"

spacy_partial_tagger/pipeline.py

@@ -3,7 +3,8 @@
 import srsly
 import torch
 from partial_tagger.data import CharBasedTags, LabelSet
-from partial_tagger.data.batch.tag import TagFactory
+from partial_tagger.data.batch.tag import TagsBatch
+from partial_tagger.data.batch.text import create_token_based_tags
 from partial_tagger.training import expected_entity_ratio_loss
 from partial_tagger.utils import create_tag
 from spacy import util
@@ -34,6 +35,8 @@ def __init__(
         self.model = model
         self.name = name
         self.scorer = scorer
+        self.padding_index = padding_index
+        self.unknown_index = unknown_index
         self.cfg: Dict[str, List[str]] = {"labels": []}
 
     @property
@@ -50,9 +53,10 @@ def set_annotations(
         tag_indices: Floats2d,
     ) -> None:
         tokenized_texts = [doc.user_data["tokenized_text"] for doc in docs]
-        tag_factory = TagFactory(tokenized_texts, self.label_set)
 
-        tags_batch = tag_factory.create_char_based_tags(tag_indices)
+        tags_batch = create_token_based_tags(
+            tokenized_texts, tag_indices, self.label_set, self.padding_index
+        )
 
         for doc, tags in zip(docs, tags_batch):
             ents = []
@@ -110,28 +114,31 @@ def get_loss(
     ) -> Tuple[float, Floats4d]:
         scores_pt = xp2torch(scores, requires_grad=True)
 
-        tokenized_texts = [
-            example.x.user_data["tokenized_text"] for example in examples
-        ]
-        tag_factory = TagFactory(tokenized_texts, self.label_set)
-
-        tags_batch = []
+        token_based_tags = []
+        lengths = []
         for example in examples:
             tags = tuple(
                 create_tag(ent.start_char, len(ent.text), ent.label_)
                 for ent in example.y.ents
             )
-            tags_batch.append(CharBasedTags(tags, example.y.text))
+            tokenized_text = example.x.user_data["tokenized_text"]
+            token_based_tags.append(
+                CharBasedTags(tags, example.x.text).convert_to_token_based(
+                    tokenized_text
+                )
+            )
+            lengths.append(tokenized_text.num_tokens)
+
+        tags_batch = TagsBatch(tuple(token_based_tags), self.label_set)
+        tags_batch.to(scores_pt.device)
+        tag_bitmap = tags_batch.get_tag_bitmap()
 
-        lengths = [text.num_tokens for text in tokenized_texts]
         max_length = max(lengths)
         mask = torch.tensor(
             [[True] * length + [False] * (max_length - length) for length in lengths],
             device=scores_pt.device,
         )
 
-        tag_bitmap = tag_factory.create_tag_bitmap(tuple(tags_batch), scores_pt.device)
-
         loss = expected_entity_ratio_loss(
             scores_pt, tag_bitmap, mask, self.label_set.get_outside_index()
         )

spacy_partial_tagger/tagger.py

@@ -3,7 +3,6 @@
 
 from partial_tagger.data import LabelSet
 from partial_tagger.data.batch.text import BaseTokenizer
-from partial_tagger.utils import create_tagger
 from spacy.tokens import Doc
 from spacy.util import registry
 from thinc.api import Model, get_torch_default_device, torch2xp, xp2torch
@@ -12,6 +11,7 @@
 from thinc.util import convert_recursive, is_torch_array, is_xp_array
 
 from .tokenizer import get_tokenizer
+from .util import create_tagger
 
 
 @registry.architectures.register("spacy-partial-tagger.PartialTagger.v1")
@@ -51,9 +51,10 @@ def forward(
         doc.user_data["tokenized_text"] = text
 
     device = get_torch_default_device()
+    text_batch.to(device)
 
     (log_potentials, tag_indices), backward = model.layers[0](
-        [text_batch.get_tagger_inputs(device), text_batch.get_mask(device)],
+        [text_batch.tagger_inputs, text_batch.mask],
         is_train,
     )
 

spacy_partial_tagger/tokenizer.py

@@ -1,11 +1,10 @@
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 from partial_tagger.data import Span, TokenizedText
 from partial_tagger.data.batch.text import (
     BaseTokenizer,
     TextBatch,
-    Texts,
     TransformerTokenizer,
 )
 from transformers import AutoTokenizer
@@ -31,7 +30,7 @@ def __init__(
         }
         self.__tokenizer_args["return_offsets_mapping"] = True
 
-    def __call__(self, texts: Texts) -> TextBatch:
+    def __call__(self, texts: Tuple[str]) -> TextBatch:
         batch_encoding = self.__tokenizer(texts, **self.__tokenizer_args)
 
         pad_token_id = self.__tokenizer.pad_token_id

spacy_partial_tagger/util.py

@@ -1,13 +1,27 @@
 from typing import List, Tuple
 
-import catalogue
 import spacy_alignments as tokenizations
-from spacy.util import registry
+from partial_tagger.data import LabelSet
+from partial_tagger.decoders.viterbi import Contrainer, ViterbiDecoder
+from partial_tagger.encoders.transformer import TransformerModelEncoderFactory
+from partial_tagger.tagger import SequenceTagger
 from transformers import PreTrainedTokenizer
 
-registry.label_indexers = catalogue.create(  # type:ignore
-    "spacy", "label_indexers", entry_points=True
-)
+
+def create_tagger(
+    model_name: str, label_set: LabelSet, padding_index: int
+) -> SequenceTagger:
+    return SequenceTagger(
+        TransformerModelEncoderFactory(model_name).create(label_set),
+        ViterbiDecoder(
+            padding_index,
+            Contrainer(
+                label_set.get_start_states(),
+                label_set.get_end_states(),
+                label_set.get_transitions(),
+            ),
+        ),
+    )
 
 
 def get_alignments(