Dependency Update (#40)

* Update dependencies

* Update a pytest setting

* Use sequence-label

Author: yasufumy

.github/workflows/ci.yml

@@ -34,4 +34,4 @@ jobs:
         mypy tests
     - name: Run tests
       run: |
-        pytest --cov=spacy_partial_tagger --cov-report=term-missing
+        pytest

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 [project]
 name = "spacy-partial-tagger"
 description = "Sequence Tagger for Partially Annotated Dataset in spaCy"
-requires-python = ">=3.8,<4.0"
+requires-python = ">=3.8,<3.12"
 readme = {file = "README.md", content-type = "text/markdown"}
 license = {file = "LICENSE"}
 authors = [
@@ -22,7 +22,8 @@ dependencies = [
     "torch<3.0.0,>=2.0.1",
     "spacy[transformers]<4.0.0,>=3.3.1",
     "spacy-alignments<1.0.0,>=0.8.5",
-    "pytorch-partial-tagger<1.0.0,>=0.1.14",
+    "pytorch-partial-tagger<1.0.0,>=0.1.15",
+    "sequence-label<1.0.0,>=0.1.4",
 ]
 dynamic = ["version"]
 
@@ -89,7 +90,4 @@ max-complexity = 18
 testpaths = [
     "tests",
 ]
-addopts = "--strict-markers -m 'not local'"
-markers = [
-    "local"
-]
+addopts = "--cov=spacy_partial_tagger --cov-report=term-missing -vv"

spacy_partial_tagger/collator.py

@@ -1,7 +1,7 @@
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
-from partial_tagger.data import Alignment, Alignments, Span
 from partial_tagger.data.collators import BaseCollator, Batch, TransformerCollator
+from sequence_label.core import LabelAlignment, Span
 from transformers import AutoTokenizer
 from transformers.models.bert_japanese import BertJapaneseTokenizer
 
@@ -23,34 +23,32 @@ def __init__(
         }
         self.__tokenizer_args["return_offsets_mapping"] = True
 
-    def __call__(self, texts: Tuple[str]) -> Tuple[Batch, Alignments]:
+    def __call__(
+        self, texts: Tuple[str, ...]
+    ) -> Tuple[Batch, Tuple[LabelAlignment, ...]]:
         batch_encoding = self.__tokenizer(texts, **self.__tokenizer_args)
 
         pad_token_id = self.__tokenizer.pad_token_id
         mask = batch_encoding.input_ids != pad_token_id
-        tokenized_text_lengths = mask.sum(dim=1)
 
         alignments = []
-        for _tokenized_text_length, input_ids, text in zip(
-            tokenized_text_lengths, batch_encoding.input_ids, texts
-        ):
+        for input_ids, text in zip(batch_encoding.input_ids, texts):
             char_spans = tuple(
-                Span(span[0], len(span)) if span else None
+                Span(start=span[0], length=len(span)) if span else None
                 for span in get_alignments(self.__tokenizer, text, input_ids.tolist())
             )
-            token_indices = [-1] * len(text)
+            token_spans: List[Optional[Span]] = [None] * len(text)
             for token_index, char_span in enumerate(char_spans):
                 if char_span is None:
                     continue
                 start = char_span.start
                 end = char_span.start + char_span.length
-                token_indices[start:end] = [token_index] * char_span.length
+                for i in range(start, end):
+                    token_spans[i] = Span(start=token_index, length=1)
 
-            alignments.append(Alignment(text, char_spans, tuple(token_indices)))
+            alignments.append(LabelAlignment(char_spans, tuple(token_spans)))
 
-        return Batch(tagger_inputs=batch_encoding, mask=mask), Alignments(
-            tuple(alignments)
-        )
+        return Batch(tagger_inputs=batch_encoding, mask=mask), tuple(alignments)
 
 
 def get_collator(

spacy_partial_tagger/pipeline.py

@@ -1,10 +1,9 @@
-from typing import Callable, Dict, Iterable, List, Optional, Set, Tuple, cast
+from typing import Callable, Dict, Iterable, List, Optional, Tuple, cast
 
 import srsly
 import torch
-from partial_tagger.data import Alignments, LabelSet
-from partial_tagger.training import compute_partially_supervised_loss
-from partial_tagger.utils import create_tag
+from partial_tagger.training import compute_partially_supervised_loss, create_tag_bitmap
+from sequence_label import LabelSet, SequenceLabel
 from spacy import util
 from spacy.errors import Errors
 from spacy.language import Language
@@ -39,7 +38,9 @@ def __init__(
 
     @property
     def label_set(self) -> LabelSet:
-        return LabelSet(set(self.cfg["labels"]))
+        return LabelSet(
+            labels=set(self.cfg["labels"]), padding_index=self.padding_index
+        )
 
     def predict(self, docs: List[Doc]) -> Floats2d:
         (_, tag_indices) = self.model.predict(docs)
@@ -50,16 +51,14 @@ def set_annotations(
         docs: List[Doc],
         tag_indices: Floats2d,
     ) -> None:
-        alignments = Alignments(tuple(doc.user_data["alignment"] for doc in docs))
-        tags_batch = alignments.create_char_based_tags(
-            tag_indices.tolist(),
-            label_set=self.label_set,
-            padding_index=self.padding_index,
+        labels = self.label_set.decode(
+            tag_indices=tag_indices.tolist(),
+            alignments=tuple(doc.user_data["alignment"] for doc in docs),
         )
 
-        for doc, tags in zip(docs, tags_batch):
+        for doc, label in zip(docs, labels):
             ents = []
-            for tag in tags:
+            for tag in label.tags:
                 span = doc.char_span(tag.start, tag.start + tag.length, tag.label)
                 if span:
                     ents.append(span)
@@ -89,19 +88,17 @@ def update(
         losses[self.name] += loss
         return losses
 
-    def initialize(
-        self, get_examples: Callable, *, nlp: Language, labels: Optional[dict] = None
-    ) -> None:
+    def initialize(self, get_examples: Callable, *, nlp: Language) -> None:
         X_small: List[Doc] = []
-        label: Set[str] = set()
+        labels: List[str] = []
         for example in get_examples():
             if len(X_small) < 10:
                 X_small.append(example.x)
             for entity in example.y.ents:
-                if entity.label_ not in label:
-                    label.add(entity.label_)
+                if entity.label_ not in labels:
+                    labels.append(entity.label_)
 
-        self.cfg["labels"] = list(label)
+        self.cfg["labels"] = labels
 
         self.model.initialize(
             X=X_small,
@@ -113,23 +110,32 @@ def get_loss(
     ) -> Tuple[float, Floats4d]:
         scores_pt = xp2torch(scores, requires_grad=True)
 
-        char_based_tags = []
-        temp = []
+        labels = []
+        alignments = []
         lengths = []
         for example in examples:
-            tags = tuple(
-                create_tag(ent.start_char, len(ent.text), ent.label_)
-                for ent in example.y.ents
+            labels.append(
+                SequenceLabel.from_dict(
+                    tags=[
+                        {
+                            "start": ent.start_char,
+                            "end": ent.end_char,
+                            "label": ent.label_,
+                        }
+                        for ent in example.y.ents
+                    ],
+                    size=len(example.y.text),
+                )
             )
-            char_based_tags.append(tags)
 
             alignment = example.x.user_data["alignment"]
-            lengths.append(alignment.num_tokens)
-            temp.append(alignment)
+            alignments.append(alignment)
+            lengths.append(alignment.target_size)
 
-        alignments = Alignments(tuple(temp))
-        tag_bitmap = torch.tensor(
-            alignments.get_tag_bitmap(char_based_tags, self.label_set),
+        tag_bitmap = create_tag_bitmap(
+            label_set=self.label_set,
+            labels=tuple(labels),
+            alignments=tuple(alignments),
             device=scores_pt.device,
         )
 
@@ -140,7 +146,7 @@ def get_loss(
         )
 
         loss = compute_partially_supervised_loss(
-            scores_pt, tag_bitmap, mask, self.label_set.get_outside_index()
+            scores_pt, tag_bitmap, mask, self.label_set.outside_index
         )
 
         (grad,) = torch.autograd.grad(loss, scores_pt)

spacy_partial_tagger/tagger.py

@@ -1,7 +1,7 @@
 from functools import partial
 from typing import Any, Callable, List, Optional, Tuple, cast
 
-from partial_tagger.data import LabelSet
+from sequence_label import LabelSet
 from spacy.tokens import Doc
 from spacy.util import registry
 from thinc.api import Model, get_torch_default_device, torch2xp, xp2torch
@@ -45,7 +45,7 @@ def forward(
     collator = model.attrs["collator"]
     batch, alignments = collator(tuple(doc.text for doc in X))
 
-    for doc, alignment in zip(X, alignments.alignments):
+    for doc, alignment in zip(X, alignments):
         doc.user_data["alignment"] = alignment
 
     device = get_torch_default_device()

spacy_partial_tagger/util.py

@@ -1,10 +1,10 @@
 from typing import List, Tuple
 
 import spacy_alignments as tokenizations
-from partial_tagger.data import LabelSet
 from partial_tagger.decoders.viterbi import Constrainer, ViterbiDecoder
 from partial_tagger.encoders.transformer import TransformerModelEncoderFactory
 from partial_tagger.tagger import SequenceTagger
+from sequence_label import LabelSet
 from transformers import PreTrainedTokenizer
 
 
@@ -16,9 +16,9 @@ def create_tagger(
         ViterbiDecoder(
             padding_index,
             Constrainer(
-                label_set.get_start_states(),
-                label_set.get_end_states(),
-                label_set.get_transitions(),
+                label_set.start_states,
+                label_set.end_states,
+                label_set.transitions,
             ),
         ),
     )

tests/test_tagger.py

@@ -1,4 +1,4 @@
-from partial_tagger.data import LabelSet
+from sequence_label import LabelSet
 from spacy.language import Language
 from spacy.tokens import Doc
 
@@ -22,6 +22,6 @@ def test_partial_tagger(nlp: Language) -> None:
     (log_potentials, tag_indices), _ = tagger(docs, is_train=False)
 
     # 10 is the length of sub-words of text.
-    num_tags = label_set.get_tag_size()
+    num_tags = label_set.state_size
     assert log_potentials.shape == (1, 10, num_tags, num_tags)
     assert tag_indices.shape == (1, 10)