iob2labels 0.1.0__tar.gz

iob2labels-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,25 @@
+# .github/workflows/publish.yml
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  id-token: write
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+
+      - name: Build package
+        run: uv build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

iob2labels-0.1.0/.gitignore

@@ -0,0 +1,11 @@
+.venv
+.env
+.DS_Store
+pyrightconfig.json
+
+__pycache__
+.pytest_cache
+
+data/
+dist/
+*.egg-info/

iob2labels-0.1.0/.python-version

@@ -0,0 +1 @@
+3.10

iob2labels-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

iob2labels-0.1.0/PKG-INFO

@@ -0,0 +1,163 @@
+Metadata-Version: 2.4
+Name: iob2labels
+Version: 0.1.0
+Summary: Convert IOB2 NER span annotations into integer label sequences aligned to any HuggingFace-compatible tokenizer.
+License-Expression: MIT
+License-File: LICENSE
+Keywords: iob2,named-entity-recognition,ner,nlp,token-classification,tokenizers
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Text Processing :: Linguistic
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Requires-Dist: pydantic>=2.0
+Requires-Dist: tokenizers>=0.15
+Description-Content-Type: text/markdown
+
+# iob2labels
+
+Convert [IOB2-format](https://en.wikipedia.org/wiki/Inside%E2%80%93outside%E2%80%93beginning_(tagging)) NER span annotations into integer label sequences for Transformer-based token classification tasks.
+
+If you use annotation tools like [Prodigy](https://prodi.gy/docs), [Label Studio](https://labelstud.io/), or [Doccano](https://github.com/doccano/doccano) to annotate NER data, this library converts those character-offset span annotations into the label arrays you need for training.
+
+## Installation
+
+```bash
+uv add iob2labels
+```
+
+Dependencies: `tokenizers` (HuggingFace Rust-backed tokenizer) and `pydantic`. No `torch` or `transformers` required.
+
+## Quick Start
+
+```python
+from iob2labels import IOB2Encoder
+
+encoder = IOB2Encoder(
+    labels=["actor", "character", "plot"],
+    tokenizer="bert-base-uncased",
+)
+
+labels = encoder(
+    text="Did Dame Judy Dench star in a British film about Queen Elizabeth?",
+    spans=[
+        {"label": "actor", "start": 4, "end": 19},
+        {"label": "plot", "start": 30, "end": 37},
+        {"label": "character", "start": 49, "end": 64},
+    ]
+)
+# >>> [-100, 0, 1, 2, 2, 2, 0, 0, 0, 5, 0, 0, 3, 4, 0, -100]
+```
+
+> Example pulled from the [MITMovie](https://groups.csail.mit.edu/sls/downloads/movie/) dataset.
+
+The output is a `list[int]` aligned to the tokenizer's output. Convert to a tensor or array as needed:
+
+```python
+import torch
+x = torch.tensor(labels)
+
+# or with numpy
+import numpy as np
+x = np.array(labels)
+```
+
+## How It Works
+
+The IOB2 format assigns each token one of three tag types:
+
+- **O** (Outside) - not part of any entity
+- **B-LABEL** (Beginning) - first token of an entity
+- **I-LABEL** (Inside) - continuation of an entity
+
+Each entity class generates 2 labels (B + I), plus the O class, so the total label count is always `(n * 2) + 1`:
+
+```python
+encoder.label_map
+# >>> {'O': 0, 'B-ACTOR': 1, 'I-ACTOR': 2, 'B-CHARACTER': 3, 'I-CHARACTER': 4, 'B-PLOT': 5, 'I-PLOT': 6}
+```
+
+Special tokens (e.g., `[CLS]`, `[SEP]`) receive the ignore value `-100`, which PyTorch's `CrossEntropyLoss` skips by default.
+
+## Tokenizer Input
+
+The `tokenizer` argument accepts three forms:
+
+```python
+# 1. checkpoint name (downloads from HuggingFace Hub)
+encoder = IOB2Encoder(labels=labels, tokenizer="bert-base-uncased")
+
+# 2. standalone tokenizers.Tokenizer instance
+from tokenizers import Tokenizer
+tok = Tokenizer.from_pretrained("bert-base-uncased")
+encoder = IOB2Encoder(labels=labels, tokenizer=tok)
+
+# 3. transformers PreTrainedTokenizerFast (unwrapped automatically)
+from transformers import AutoTokenizer
+tok = AutoTokenizer.from_pretrained("bert-base-uncased")
+encoder = IOB2Encoder(labels=labels, tokenizer=tok)
+```
+
+## Batch Encoding
+
+```python
+annotations = [
+    {"text": "Did Dame Judy Dench star?", "spans": [{"label": "actor", "start": 4, "end": 19}]},
+    {"text": "Matt Damon was Jason Bourne.", "spans": [{"label": "actor", "start": 0, "end": 10}]},
+]
+
+results = encoder.batch(annotations)
+# >>> [[-100, 0, 1, 2, 2, 2, 0, -100], [-100, 1, 2, 0, 0, 0, 0, -100]]
+```
+
+The batch path uses the Rust-backed `encode_batch()` for parallelized tokenization. Returns `list[list[int]]` with no padding; use HuggingFace's `DataCollatorForTokenClassification` or your own padding for training.
+
+## Custom Field Names
+
+If your annotation data uses non-standard field names, configure them at construction:
+
+```python
+# BioMed-NER dataset uses "entities" and "class" instead of "spans" and "label"
+encoder = IOB2Encoder(
+    labels=["organism", "chemicals"],
+    tokenizer="bert-base-uncased",
+    spans_field="entities",
+    label_field="class",
+)
+```
+
+## Built-in Conversion Check
+
+By default, every encoding is verified by recovering the entity text from the produced labels and comparing it to the original annotation. This catches misalignment bugs early. Disable it for performance in production:
+
+```python
+encoder = IOB2Encoder(labels=labels, tokenizer=tok, conversion_check=False)
+```
+
+## Supported Tokenizers
+
+Tested across three tokenizer families:
+
+| Family | Checkpoints |
+|---|---|
+| WordPiece | `bert-base-cased`, `bert-base-uncased`, `bert-large-cased`, `bert-large-uncased`, `distilbert-base-cased`, `distilbert-base-uncased`, `google/electra-base-discriminator` |
+| BPE | `roberta-base`, `roberta-large` |
+| SentencePiece | `albert-base-v2`, `xlnet-base-cased`, `t5-small` |
+
+Other HuggingFace-compatible tokenizers should work as well. The built-in conversion check will flag any issues.
+
+## Tests
+
+```bash
+uv run pytest tests/ -v
+```
+
+The test suite includes unit tests for label map construction, entity range detection, and the conversion checker, plus a parametrized matrix of 12 tokenizer checkpoints across multiple annotation edge cases (entities at text boundaries, adjacent entities, punctuation, etc.).

iob2labels-0.1.0/README.md

@@ -0,0 +1,139 @@
+# iob2labels
+
+Convert [IOB2-format](https://en.wikipedia.org/wiki/Inside%E2%80%93outside%E2%80%93beginning_(tagging)) NER span annotations into integer label sequences for Transformer-based token classification tasks.
+
+If you use annotation tools like [Prodigy](https://prodi.gy/docs), [Label Studio](https://labelstud.io/), or [Doccano](https://github.com/doccano/doccano) to annotate NER data, this library converts those character-offset span annotations into the label arrays you need for training.
+
+## Installation
+
+```bash
+uv add iob2labels
+```
+
+Dependencies: `tokenizers` (HuggingFace Rust-backed tokenizer) and `pydantic`. No `torch` or `transformers` required.
+
+## Quick Start
+
+```python
+from iob2labels import IOB2Encoder
+
+encoder = IOB2Encoder(
+    labels=["actor", "character", "plot"],
+    tokenizer="bert-base-uncased",
+)
+
+labels = encoder(
+    text="Did Dame Judy Dench star in a British film about Queen Elizabeth?",
+    spans=[
+        {"label": "actor", "start": 4, "end": 19},
+        {"label": "plot", "start": 30, "end": 37},
+        {"label": "character", "start": 49, "end": 64},
+    ]
+)
+# >>> [-100, 0, 1, 2, 2, 2, 0, 0, 0, 5, 0, 0, 3, 4, 0, -100]
+```
+
+> Example pulled from the [MITMovie](https://groups.csail.mit.edu/sls/downloads/movie/) dataset.
+
+The output is a `list[int]` aligned to the tokenizer's output. Convert to a tensor or array as needed:
+
+```python
+import torch
+x = torch.tensor(labels)
+
+# or with numpy
+import numpy as np
+x = np.array(labels)
+```
+
+## How It Works
+
+The IOB2 format assigns each token one of three tag types:
+
+- **O** (Outside) - not part of any entity
+- **B-LABEL** (Beginning) - first token of an entity
+- **I-LABEL** (Inside) - continuation of an entity
+
+Each entity class generates 2 labels (B + I), plus the O class, so the total label count is always `(n * 2) + 1`:
+
+```python
+encoder.label_map
+# >>> {'O': 0, 'B-ACTOR': 1, 'I-ACTOR': 2, 'B-CHARACTER': 3, 'I-CHARACTER': 4, 'B-PLOT': 5, 'I-PLOT': 6}
+```
+
+Special tokens (e.g., `[CLS]`, `[SEP]`) receive the ignore value `-100`, which PyTorch's `CrossEntropyLoss` skips by default.
+
+## Tokenizer Input
+
+The `tokenizer` argument accepts three forms:
+
+```python
+# 1. checkpoint name (downloads from HuggingFace Hub)
+encoder = IOB2Encoder(labels=labels, tokenizer="bert-base-uncased")
+
+# 2. standalone tokenizers.Tokenizer instance
+from tokenizers import Tokenizer
+tok = Tokenizer.from_pretrained("bert-base-uncased")
+encoder = IOB2Encoder(labels=labels, tokenizer=tok)
+
+# 3. transformers PreTrainedTokenizerFast (unwrapped automatically)
+from transformers import AutoTokenizer
+tok = AutoTokenizer.from_pretrained("bert-base-uncased")
+encoder = IOB2Encoder(labels=labels, tokenizer=tok)
+```
+
+## Batch Encoding
+
+```python
+annotations = [
+    {"text": "Did Dame Judy Dench star?", "spans": [{"label": "actor", "start": 4, "end": 19}]},
+    {"text": "Matt Damon was Jason Bourne.", "spans": [{"label": "actor", "start": 0, "end": 10}]},
+]
+
+results = encoder.batch(annotations)
+# >>> [[-100, 0, 1, 2, 2, 2, 0, -100], [-100, 1, 2, 0, 0, 0, 0, -100]]
+```
+
+The batch path uses the Rust-backed `encode_batch()` for parallelized tokenization. Returns `list[list[int]]` with no padding; use HuggingFace's `DataCollatorForTokenClassification` or your own padding for training.
+
+## Custom Field Names
+
+If your annotation data uses non-standard field names, configure them at construction:
+
+```python
+# BioMed-NER dataset uses "entities" and "class" instead of "spans" and "label"
+encoder = IOB2Encoder(
+    labels=["organism", "chemicals"],
+    tokenizer="bert-base-uncased",
+    spans_field="entities",
+    label_field="class",
+)
+```
+
+## Built-in Conversion Check
+
+By default, every encoding is verified by recovering the entity text from the produced labels and comparing it to the original annotation. This catches misalignment bugs early. Disable it for performance in production:
+
+```python
+encoder = IOB2Encoder(labels=labels, tokenizer=tok, conversion_check=False)
+```
+
+## Supported Tokenizers
+
+Tested across three tokenizer families:
+
+| Family | Checkpoints |
+|---|---|
+| WordPiece | `bert-base-cased`, `bert-base-uncased`, `bert-large-cased`, `bert-large-uncased`, `distilbert-base-cased`, `distilbert-base-uncased`, `google/electra-base-discriminator` |
+| BPE | `roberta-base`, `roberta-large` |
+| SentencePiece | `albert-base-v2`, `xlnet-base-cased`, `t5-small` |
+
+Other HuggingFace-compatible tokenizers should work as well. The built-in conversion check will flag any issues.
+
+## Tests
+
+```bash
+uv run pytest tests/ -v
+```
+
+The test suite includes unit tests for label map construction, entity range detection, and the conversion checker, plus a parametrized matrix of 12 tokenizer checkpoints across multiple annotation edge cases (entities at text boundaries, adjacent entities, punctuation, etc.).

iob2labels-0.1.0/pyproject.toml

@@ -0,0 +1,40 @@
+[project]
+name = "iob2labels"
+version = "0.1.0"
+description = "Convert IOB2 NER span annotations into integer label sequences aligned to any HuggingFace-compatible tokenizer."
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.10"
+keywords = ["ner", "named-entity-recognition", "iob2", "token-classification", "nlp", "tokenizers"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Text Processing :: Linguistic",
+    "Typing :: Typed",
+]
+dependencies = [
+    "pydantic>=2.0",
+    "tokenizers>=0.15",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/iob2labels"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+dev = [
+    "ipython>=8.37.0",
+    "pytest>=8.4.2",
+    "transformers>=4.40",
+]

iob2labels-0.1.0/src/iob2labels/__init__.py

@@ -0,0 +1,15 @@
+## -- primary interface
+from .encoder import IOB2Encoder as IOB2Encoder
+
+## -- types
+from .annotations import Annotation as Annotation
+from .annotations import Span as Span
+
+## -- utilities
+from .labels import create_label_map as create_label_map
+from .labels import format_entity_label as format_entity_label
+from .annotations import preprocessing as preprocessing
+
+## -- checker
+from .checker import check_iob_conversion as check_iob_conversion
+from .checker import get_entity_index_ranges as get_entity_index_ranges

iob2labels-0.1.0/src/iob2labels/annotations.py

@@ -0,0 +1,87 @@
+from typing import Any
+from typing_extensions import TypedDict
+
+from pydantic import BaseModel, Field, ValidationError
+from pydantic import AfterValidator, StrictStr, StrictInt
+
+
+class DefaultFields:
+    TEXT = "text"
+    SPANS = "spans"
+    START = "start"
+    END = "end"
+    LABEL = "label"
+
+## -- Pydantic models to perform validation of input annotation data
+## -- and support conversion for iob-label conversion.
+
+class _SpanFormat(BaseModel):
+    start: StrictInt = Field(..., description="Index of entity starting character in text string.")
+    end: StrictInt = Field(..., description="Index of entity ending character in text string.")
+    label: StrictStr = Field(..., description="Label name for annotated entity (e.g., PERSON, PRODUCT, LOCATION, etc.")
+
+class _AnnotationFormat(BaseModel):
+    text: StrictStr
+    spans: list[_SpanFormat]
+
+def convert_to_validated_format(
+    text: str,
+    spans: list[dict],
+    start_field: str,
+    end_field: str,
+    label_field: str
+) -> _AnnotationFormat:
+    """Intermediate function for converting input annotation data to pydantic models for validation and field conversion."""
+    return _AnnotationFormat(
+        text=text,
+        spans=[
+            _SpanFormat(
+                start=span[start_field],
+                end=span[end_field],
+                label=span[label_field]
+            ) for span in spans
+        ]
+    )
+
+## -- typed-dicts are returned from validation step as they are _lightly_ typed
+## -- but benefit from intermediate pydantic validation and preprocessing.
+class Span(TypedDict):
+    start: int
+    end: int
+    label: str
+
+class Annotation(TypedDict):
+    text: str
+    spans: list[Span]
+
+def preprocessing(
+    text: str,
+    spans: list[dict],
+    start_field: str = DefaultFields.START,
+    end_field: str = DefaultFields.END,
+    label_field: str = DefaultFields.LABEL
+) -> Annotation:
+    # first convert to pydantic models to convert and validate input data
+    validated = convert_to_validated_format(text, spans, start_field, end_field, label_field)
+    # return as typed dict
+    return Annotation(**validated.model_dump())
+
+def validate_batch(
+    annotations: list[dict],
+    text_field: str = DefaultFields.TEXT,
+    spans_field: str = DefaultFields.SPANS,
+    start_field: str = DefaultFields.START,
+    end_field: str = DefaultFields.END,
+    label_field: str = DefaultFields.LABEL
+) -> list[Annotation]:
+    assert isinstance(annotations, list) and all([isinstance(ann, dict) for ann in annotations]), f"Input for annotations is not a list of dicts."
+    return [
+        preprocessing(  # <- was `validate` (nonexistent); fixed to call preprocessing
+            text=ann[text_field],
+            spans=ann[spans_field],
+            start_field=start_field,
+            end_field=end_field,
+            label_field=label_field
+        )
+        for ann in annotations
+    ]

iob2labels-0.1.0/src/iob2labels/checker.py

@@ -0,0 +1,71 @@
+from itertools import takewhile
+
+from tokenizers import Tokenizer
+
+from iob2labels.labels import IobPrefixes, IGNORE_TOKEN
+from iob2labels.annotations import Annotation, DefaultFields
+
+def invert_label_map(label_map: dict[str, int]) -> dict[int, str]:
+ return {v: k for k, v in label_map.items()}
+
+def get_iob_type_by_iob_label(label_map: dict[str, int], iob_label: int) -> str:
+    if iob_label == IGNORE_TOKEN:
+        iob_type = IobPrefixes.OUTSIDE
+    else:
+        idx_map = invert_label_map(label_map)
+        iob_type = idx_map[iob_label]
+    return iob_type[0]
+
+def is_beginning_tag(label_map: dict[str, int], iob_label: int) -> bool:
+    """Boolean check if label associated with input index is a beginning tag."""
+    return get_iob_type_by_iob_label(label_map, iob_label) == IobPrefixes.BEGINNING
+
+def is_inside_tag(label_map: dict[str, int], iob_label: int) -> bool:
+    """Boolean check if label associated with input index is an inside tag."""
+    return get_iob_type_by_iob_label(label_map, iob_label) == IobPrefixes.INSIDE
+
+def is_outside_tag(label_map: dict[str, int], iob_label: int) -> bool:
+    """Boolean check if label associated with input index is an outside tag."""
+    return get_iob_type_by_iob_label(label_map, iob_label) == IobPrefixes.OUTSIDE
+
+
+
+def get_entity_sequence_length(label_map: dict[str, int], iob_labels: list[int]) -> int:
+    return len(list(takewhile(lambda x: is_inside_tag(label_map, x), iob_labels)))
+
+def get_entity_index_ranges(label_map: dict[str, int], iob_labels: list[int]) -> list[tuple[int, int]]:
+    return [
+        (idx, idx + get_entity_sequence_length(label_map, iob_labels[(idx + 1):]))
+        for idx in range(len(iob_labels)) if is_beginning_tag(label_map, iob_labels[idx])
+    ]
+
+def check_iob_conversion(
+    iob_labels: list[int],
+    label_map: dict[str, int],
+    tokenizer: Tokenizer,
+    input_ids: list[int],
+    annotation: Annotation,
+    debug: bool = False,
+    strict: bool = True,
+) -> None:
+    """Tests to ensure assigned IOB labels are correct based on character and token indices for annotated entities."""
+    match_ranges = get_entity_index_ranges(label_map, iob_labels)
+    num_ranges, num_spans = len(match_ranges), len(annotation[DefaultFields.SPANS])
+    assert num_ranges == num_spans, f"Test found {num_ranges} matches but annotation includes {num_spans} entities."
+
+    for _range, span in zip(match_ranges, annotation[DefaultFields.SPANS]):
+        # recover entity from IOB positive label indices
+        entity_input_ids = input_ids[_range[0]: (_range[1] + 1)]
+        recovered_entity = tokenizer.decode(entity_input_ids).strip() # <- some tokenizers (e.g., Roberta-Base), can include leading whitespace in decoded entity
+
+        # encode/decode annotated entity and assert equality
+        annotated_entity = annotation[DefaultFields.TEXT][span[DefaultFields.START]:span[DefaultFields.END]]
+        expected_entity = tokenizer.decode(
+            tokenizer.encode(annotated_entity, add_special_tokens=False).ids  # <- standalone tokenizers: encode() returns Encoding, access .ids for token IDs
+        )
+        result = expected_entity == recovered_entity if strict else expected_entity in recovered_entity
+        assert result, f"Recovered entity (via IOB labels) '{recovered_entity}' does not match expected entity '{annotated_entity}'. Decoded form is '{expected_entity}'."
+
+        if debug: print(f"| -> recovered entity '{recovered_entity}' for IOB labels at indices ({_range[0]}, {_range[1]}), which matches annotated entity '{annotated_entity}'.")
+
+    return