span_runtime.py

"""Keep offsets relative to original text with GLiNER2's span architecture."""
import copy


class OriginalTextCollator:
    def __init__(self, processor):
        self.processor = processor

    def __call__(self, batch):
        self.processor.change_mode(False)
        transformed = []
        for text, schema in batch:
            if not text:
                raise ValueError('empty_text_requires_an_explicit_empty_result')
            if hasattr(schema, 'build'):
                schema = schema.build()
            transformed.append(self.processor._transform_record(
                {'text': text, 'schema': copy.deepcopy(schema)}, max_len=None))
        return self.processor._pad_batch(transformed)


def preserve_original_text(model):
    """Install before using batch_extract/extract with max_len=None.

    Native GLiNER2 2.0.0's generic collator appends a period to unpunctuated
    inputs. This collator keeps both the encoder input and span offsets faithful
    to the supplied text. Long-input chunking must happen before this call.
    """
    if getattr(model, 'architecture', 'span') != 'span':
        raise ValueError('this_runtime_adapter_is_for_span_checkpoints')
    if getattr(model.config, 'source_word_splitter', None) == 'fine-source-v1':
        from fine_span_tokenizer import fine_words
        model.processor.word_splitter = fine_words
    model._inference_collator = OriginalTextCollator(model.processor)
    return model