Named Entity Recognition

Named‑Entity Recognition (NER) predictive task.

Entity

Bases: BaseModel

Entity mention with text span and type.

Source code in sieves/tasks/predictive/ner/core.py

class Entity(pydantic.BaseModel):
    """Entity mention with text span and type."""

    text: str
    context: str
    entity_type: str

FewshotExample

Bases: FewshotExample

Few‑shot example with entities annotated in text.

Source code in sieves/tasks/predictive/ner/core.py

class FewshotExample(BaseFewshotExample):
    """Few‑shot example with entities annotated in text."""

    text: str
    entities: list[Entity]

    @override
    @property
    def target_fields(self) -> Sequence[str]:
        return ("entities",)

input_fields property

Defines which fields are inputs.

Returns:

Type	Description
Sequence[str]	Sequence of field names.

from_dspy(example) classmethod

Convert from dspy.Example.

Parameters:

Name	Type	Description	Default
example	Example	Example as dspy.Example.	required

Returns:

Type	Description
Self	Example as FewshotExample.

Source code in sieves/tasks/predictive/core.py

@classmethod
def from_dspy(cls, example: dspy.Example) -> Self:
    """Convert from `dspy.Example`.

    :param example: Example as `dspy.Example`.
    :returns: Example as `FewshotExample`.
    """
    return cls(**example)

to_dspy()

Convert to dspy.Example.

Returns:

Type	Description
Example	Example as dspy.Example.

Source code in sieves/tasks/predictive/core.py

def to_dspy(self) -> dspy.Example:
    """Convert to `dspy.Example`.

    :returns: Example as `dspy.Example`.
    """
    return dspy.Example(**Engine.convert_fewshot_examples([self])[0]).with_inputs(self.input_fields)

NER

Bases: PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]

Extract named entities from text using various engines.

Source code in sieves/tasks/predictive/ner/core.py

class NER(PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]):
    """Extract named entities from text using various engines."""

    def __init__(
        self,
        entities: list[str],
        model: _TaskModel,
        task_id: str | None = None,
        include_meta: bool = True,
        batch_size: int = -1,
        prompt_instructions: str | None = None,
        fewshot_examples: Sequence[FewshotExample] = (),
        generation_settings: GenerationSettings = GenerationSettings(),
    ) -> None:
        """
        Initialize NER task.

        :param entities: List of entities to extract.
        :param model: Model to use.
        :param task_id: Task ID.
        :param include_meta: Whether to include meta information generated by the task.
        :param batch_size: Batch size to use for inference. Use -1 to process all documents at once.
        :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
        :param fewshot_examples: Few-shot examples.
        :param generation_settings: Settings for structured generation.
        """
        self._entities = entities or ["PERSON", "LOCATION", "ORGANIZATION"]
        super().__init__(
            model=model,
            task_id=task_id,
            include_meta=include_meta,
            batch_size=batch_size,
            overwrite=False,
            prompt_instructions=prompt_instructions,
            fewshot_examples=fewshot_examples,
            generation_settings=generation_settings,
        )
        self._fewshot_examples: Sequence[FewshotExample]

    @override
    def _init_bridge(self, engine_type: EngineType) -> _TaskBridge:
        bridge_types = {
            EngineType.langchain: LangChainNER,
            EngineType.outlines: OutlinesNER,
            EngineType.dspy: DSPyNER,
            EngineType.glix: GliXNER,
        }
        try:
            bridge_class = bridge_types[engine_type]
            result = bridge_class(
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                entities=self._entities,
            )
            return result  # type: ignore[return-value]
        except KeyError as err:
            raise KeyError(f"Engine type {engine_type} is not supported by {self.__class__.__name__}.") from err

    @override
    @property
    def supports(self) -> set[EngineType]:
        return {
            EngineType.langchain,
            EngineType.dspy,
            EngineType.outlines,
            EngineType.glix,
        }

    @override
    def _validate_fewshot_examples(self) -> None:
        for fs_example in self._fewshot_examples or []:
            for entity in fs_example.entities:
                if entity.entity_type not in self._entities:
                    raise ValueError(f"Entity {entity.entity_type} not in {self._entities}.")

    @override
    @property
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "entities": self._entities,
        }

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float = 0.5) -> datasets.Dataset:
        # Define metadata and features for the dataset
        features = datasets.Features(
            {
                "text": datasets.Value("string"),
                "entities": datasets.Sequence(
                    datasets.Features(
                        {
                            "text": datasets.Value("string"),
                            "start": datasets.Value("int32"),
                            "end": datasets.Value("int32"),
                            "entity_type": datasets.Value("string"),
                        }
                    )
                ),
            }
        )

        info = datasets.DatasetInfo(
            description=f"Named Entity Recognition dataset with entity types {self._entities}. Generated with sieves "
            f"v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset
        try:
            data: list[tuple[str, list[dict[str, Any]]]] = []
            for doc in docs:
                if self._task_id not in doc.results:
                    raise KeyError(f"Document does not have results for task ID {self._task_id}")

                # Get the entities from the document results
                result = doc.results[self._task_id].entities
                entities: list[dict[str, Any]] = []

                # List format (could be list of dictionaries or other entities)
                for entity in result:
                    assert hasattr(entity, "text")
                    assert hasattr(entity, "start")
                    assert hasattr(entity, "end")
                    assert hasattr(entity, "entity_type")

                    entities.append(
                        {
                            "text": entity.text,
                            "start": entity.start,
                            "end": entity.end,
                            "entity_type": entity.entity_type,
                        }
                    )

                data.append((doc.text or "", entities))

        except KeyError as err:
            raise KeyError(f"Not all documents have results for this task with ID {self._task_id}") from err

        def generate_data() -> Iterable[dict[str, Any]]:
            """Yield results as dicts.

            :return: Results as dicts.
            """
            for text, entities in data:
                yield {"text": text, "entities": entities}

        # Create dataset
        return datasets.Dataset.from_generator(generate_data, features=features, info=info)

    @override
    def distill(
        self,
        base_model_id: str,
        framework: DistillationFramework,
        data: datasets.Dataset | Sequence[Doc],
        output_path: Path | str,
        val_frac: float,
        init_kwargs: dict[str, Any] | None = None,
        train_kwargs: dict[str, Any] | None = None,
        seed: int | None = None,
    ) -> None:
        raise NotImplementedError

    @override
    def _evaluate_optimization_example(
        self, truth: dspy.Example, pred: dspy.Prediction, model: dspy.LM, trace: Any | None = None
    ) -> float:
        # Compute entity-level F1 score based on (text, entity_type) pairs
        true_entities = {(e["text"], e["entity_type"]) for e in truth["entities"]}
        pred_entities = {(e["text"], e["entity_type"]) for e in pred.get("entities", [])}

        if not true_entities:
            return 1.0 if not pred_entities else 0.0

        precision = len(true_entities & pred_entities) / len(pred_entities) if pred_entities else 0
        recall = len(true_entities & pred_entities) / len(true_entities)
        return 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0

fewshot_examples property

Return few-shot examples.

Returns:

Type	Description
Sequence[FewshotExample]	Few-shot examples.

id property

Return task ID.

Used by pipeline for results and dependency management.

Returns:

Type	Description
str	Task ID.

prompt_signature_description property

Return prompt signature description.

Returns:

Type	Description
str | None	Prompt signature description.

prompt_template property

Return prompt template.

Returns:

Type	Description
str	Prompt template.

__add__(other)

Chain this task with another task or pipeline using the + operator.

This returns a new Pipeline that executes this task first, followed by the task(s) in other. The original task(s)/pipeline are not mutated.

Cache semantics: - If other is a Pipeline, the resulting pipeline adopts other's use_cache setting (because the left-hand side is a single task). - If other is a Task, the resulting pipeline defaults to use_cache=True.

Parameters:

Name	Type	Description	Default
other	Task | Pipeline	A Task or Pipeline to execute after this task.	required

Returns:

Type	Description
Pipeline	A new Pipeline representing the chained execution.

Raises:

Type	Description
TypeError	If other is not a Task or Pipeline.

Source code in sieves/tasks/core.py

def __add__(self, other: Task | Pipeline) -> Pipeline:
    """Chain this task with another task or pipeline using the ``+`` operator.

    This returns a new ``Pipeline`` that executes this task first, followed by the
    task(s) in ``other``. The original task(s)/pipeline are not mutated.

    Cache semantics:
    - If ``other`` is a ``Pipeline``, the resulting pipeline adopts ``other``'s
      ``use_cache`` setting (because the left-hand side is a single task).
    - If ``other`` is a ``Task``, the resulting pipeline defaults to ``use_cache=True``.

    :param other: A ``Task`` or ``Pipeline`` to execute after this task.
    :return: A new ``Pipeline`` representing the chained execution.
    :raises TypeError: If ``other`` is not a ``Task`` or ``Pipeline``.
    """
    # Lazy import to avoid circular dependency at module import time.
    from sieves.pipeline import Pipeline

    if isinstance(other, Pipeline):
        return Pipeline(tasks=[self, *other.tasks], use_cache=other.use_cache)

    if isinstance(other, Task):
        return Pipeline(tasks=[self, other])

    raise TypeError(f"Cannot chain Task with {type(other).__name__}")

__call__(docs)

Execute the task on a set of documents.

Parameters:

Name	Type	Description	Default
docs	Iterable[Doc]	Documents to process.	required

Returns:

Type	Description
Iterable[Doc]	Processed documents.

Source code in sieves/tasks/predictive/core.py

def __call__(self, docs: Iterable[Doc]) -> Iterable[Doc]:
    """Execute the task on a set of documents.

    :param docs: Documents to process.
    :return Iterable[Doc]: Processed documents.
    """
    # 1. Compile expected prompt signatures.
    signature = self._bridge.prompt_signature

    # 2. Build executable.
    executable = self._engine.build_executable(
        inference_mode=self._bridge.inference_mode,
        prompt_template=self.prompt_template,
        prompt_signature=signature,
        fewshot_examples=self._fewshot_examples,
    )

    # Compute batch-wise results.
    batch_size = self._batch_size if self._batch_size > 0 else sys.maxsize
    while docs_batch := [doc for doc in itertools.islice(docs, batch_size)]:
        if len(docs_batch) == 0:
            break

        # 3. Extract values from docs to inject/render those into prompt templates.
        docs_values = list(self._bridge.extract(docs_batch))
        assert len(docs_values) == len(docs_batch)

        # 4. Map extracted docs values onto chunks.
        docs_chunks_offsets: list[tuple[int, int]] = []
        docs_chunks: list[dict[str, Any]] = []
        for doc, doc_values in zip(docs_batch, docs_values):
            assert doc.text
            doc_chunks_values = [doc_values | {"text": chunk} for chunk in (doc.chunks or [doc.text])]
            docs_chunks_offsets.append((len(docs_chunks), len(docs_chunks) + len(doc_chunks_values)))
            docs_chunks.extend(doc_chunks_values)

        # 5. Execute prompts per chunk.
        results = list(executable(docs_chunks))
        assert len(results) == len(docs_chunks)

        # 6. Consolidate chunk results.
        results = list(self._bridge.consolidate(results, docs_chunks_offsets))
        assert len(results) == len(docs_batch)

        # 7. Integrate results into docs.
        docs_batch = self._bridge.integrate(results, docs_batch)

        yield from docs_batch

__init__(entities, model, task_id=None, include_meta=True, batch_size=-1, prompt_instructions=None, fewshot_examples=(), generation_settings=GenerationSettings())

Initialize NER task.

Parameters:

Name	Type	Description	Default
entities	list[str]	List of entities to extract.	required
model	_TaskModel	Model to use.	required
task_id	str | None	Task ID.	None
include_meta	bool	Whether to include meta information generated by the task.	True
batch_size	int	Batch size to use for inference. Use -1 to process all documents at once.	-1
prompt_instructions	str | None	Custom prompt instructions. If None, default instructions are used.	None
fewshot_examples	Sequence[FewshotExample]	Few-shot examples.	()
generation_settings	GenerationSettings	Settings for structured generation.	GenerationSettings()

Source code in sieves/tasks/predictive/ner/core.py

def __init__(
    self,
    entities: list[str],
    model: _TaskModel,
    task_id: str | None = None,
    include_meta: bool = True,
    batch_size: int = -1,
    prompt_instructions: str | None = None,
    fewshot_examples: Sequence[FewshotExample] = (),
    generation_settings: GenerationSettings = GenerationSettings(),
) -> None:
    """
    Initialize NER task.

    :param entities: List of entities to extract.
    :param model: Model to use.
    :param task_id: Task ID.
    :param include_meta: Whether to include meta information generated by the task.
    :param batch_size: Batch size to use for inference. Use -1 to process all documents at once.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    :param fewshot_examples: Few-shot examples.
    :param generation_settings: Settings for structured generation.
    """
    self._entities = entities or ["PERSON", "LOCATION", "ORGANIZATION"]
    super().__init__(
        model=model,
        task_id=task_id,
        include_meta=include_meta,
        batch_size=batch_size,
        overwrite=False,
        prompt_instructions=prompt_instructions,
        fewshot_examples=fewshot_examples,
        generation_settings=generation_settings,
    )
    self._fewshot_examples: Sequence[FewshotExample]

deserialize(config, **kwargs) classmethod

Generate PredictiveTask instance from config.

Parameters:

Name	Type	Description	Default
config	Config	Config to generate instance from.	required
kwargs	dict[str, Any]	Values to inject into loaded config.	{}

Returns:

Type	Description
PredictiveTask[TaskPromptSignature, TaskResult, TaskBridge]	Deserialized PredictiveTask instance.

Source code in sieves/tasks/predictive/core.py

@classmethod
def deserialize(
    cls, config: Config, **kwargs: dict[str, Any]
) -> PredictiveTask[TaskPromptSignature, TaskResult, TaskBridge]:
    """Generate PredictiveTask instance from config.

    :param config: Config to generate instance from.
    :param kwargs: Values to inject into loaded config.
    :return PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]: Deserialized PredictiveTask instance.
    """
    init_dict = config.to_init_dict(cls, **kwargs)
    init_dict["generation_settings"] = GenerationSettings.model_validate(init_dict["generation_settings"])

    return cls(**init_dict)

optimize(optimizer, verbose=True)

Optimize task prompt and few-shot examples with the available optimization config.

Updates task to use best prompt and few-shot examples found by the optimizer.

Parameters:

Name	Type	Description	Default
optimizer	Optimizer	Optimizer to run.	required
verbose	bool	Whether to suppress output. DSPy produces a good amount of logs, so this can be useful to not pollute your terminal. Only warnings and errors will be printed.	True

Returns:

Type	Description
tuple[str, Sequence[FewshotExample]]	Best found prompt and few-shot examples.

Source code in sieves/tasks/predictive/core.py

def optimize(self, optimizer: optimization.Optimizer, verbose: bool = True) -> tuple[str, Sequence[FewshotExample]]:
    """Optimize task prompt and few-shot examples with the available optimization config.

    Updates task to use best prompt and few-shot examples found by the optimizer.

    :param optimizer: Optimizer to run.
    :param verbose: Whether to suppress output. DSPy produces a good amount of logs, so this can be useful to
        not pollute your terminal. Only warnings and errors will be printed.

    :return tuple[str, Sequence[FewshotExample]]: Best found prompt and few-shot examples.
    """
    assert len(self._fewshot_examples) > 1, "At least two few-shot examples need to be provided to optimize."

    # Run optimizer to get best prompt and few-shot examples.
    signature = self._get_task_signature()
    dspy_examples = [ex.to_dspy() for ex in self._fewshot_examples]
    pred_eval = functools.partial(self._evaluate_optimization_example, model=optimizer.model)

    if verbose:
        best_prompt, best_examples = optimizer(signature, dspy_examples, pred_eval, verbose=verbose)
    else:
        # Temporarily suppress DSPy logs.
        dspy_logger = logging.getLogger("dspy")
        optuna_logger = logging.getLogger("optuna")
        original_dspy_level = dspy_logger.level
        original_optuna_level = optuna_logger.level

        try:
            dspy_logger.setLevel(logging.ERROR)
            optuna_logger.setLevel(logging.ERROR)
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                best_prompt, best_examples = optimizer(signature, dspy_examples, pred_eval, verbose=verbose)
        finally:
            dspy_logger.setLevel(original_dspy_level)
            optuna_logger.setLevel(original_optuna_level)

    # Update few-shot examples and prompt instructions.
    fewshot_example_cls = self._fewshot_examples[0].__class__
    self._fewshot_examples = [fewshot_example_cls.from_dspy(ex) for ex in best_examples]
    self._validate_fewshot_examples()
    self._custom_prompt_instructions = best_prompt

    # Reinitialize bridge to use new prompt and few-shot examples.
    self._bridge = self._init_bridge(EngineType.get_engine_type(self._engine))

    return best_prompt, self._fewshot_examples

serialize()

Serialize task.

Returns:

Type	Description
Config	Config instance.

Source code in sieves/tasks/core.py

def serialize(self) -> Config:
    """Serialize task.

    :return: Config instance.
    """
    return Config.create(self.__class__, {k: Attribute(value=v) for k, v in self._state.items()})

Bridges for NER task.

DSPyNER

Bases: NERBridge[PromptSignature, Result, InferenceMode]

DSPy bridge for NER.

Source code in sieves/tasks/predictive/ner/bridges.py

class DSPyNER(NERBridge[dspy_.PromptSignature, dspy_.Result, dspy_.InferenceMode]):
    """DSPy bridge for NER."""

    @override
    @property
    def _default_prompt_instructions(self) -> str:
        return """
        A named entity recognition result that represents named entities from the provided text.
        For each entity found it includes:
        - exact text of the entity
        - a context string that contains the exact entity text along with a few surrounding words
          (two or three surronding words). The context includes the entity text itself.
        - if the same entity appears multiple times in the text, each occurrence is listed separately with its
        own context
        - the entity type from the provided list of entity types. Only entities of the specified types are included.
        """

    @override
    @property
    def _prompt_example_template(self) -> str | None:
        return None

    @override
    @property
    def _prompt_conclusion(self) -> str | None:
        return None

    @override
    @cached_property
    def prompt_signature(self) -> type[dspy_.PromptSignature]:
        entity_types = self._entities
        LiteralType = Literal[*entity_types]  # type: ignore[valid-type]

        class Entity(dspy.Signature):  # type: ignore[misc]
            text: str = dspy.OutputField(
                description="The extracted entity text, if the same entity appears multiple times in the text, "
                "includes each occurrence separately."
            )
            context: str = dspy.OutputField(
                description="A context string that MUST include the exact entity text. The context should include "
                "the entity and a few surrounding words (two or three surrounding words). IMPORTANT: The entity text "
                "MUST be present in the context string exactly as it appears in the text."
            )
            entity_type: LiteralType = dspy.OutputField(description="The type of entity")

        class Prediction(dspy.Signature):  # type: ignore[misc]
            text: str = dspy.InputField(description="Text to extract entities from")
            entity_types: list[str] = dspy.InputField(description="List of entity types to extract")

            entities: list[Entity] = dspy.OutputField(
                description="List of entities found in the text. If the same entity appears multiple times "
                "in different contexts, include each occurrence separately."
            )

        Prediction.__doc__ = jinja2.Template(self._prompt_instructions).render()

        return Prediction

    @override
    @property
    def inference_mode(self) -> dspy_.InferenceMode:
        return dspy_.InferenceMode.predict

    @override
    def consolidate(
        self, results: Iterable[dspy_.Result], docs_offsets: list[tuple[int, int]]
    ) -> Iterable[dspy_.Result]:
        results = list(results)
        # Process each document (which may consist of multiple chunks)
        for doc_offset in docs_offsets:
            doc_results = results[doc_offset[0] : doc_offset[1]]

            # Combine all entities from all chunks
            all_entities: list[Entity] = []

            # Process each chunk for this document
            for chunk_result in doc_results:
                if not hasattr(chunk_result, "entities") or not chunk_result.entities:
                    continue

                # Process entities in this chunk
                for entity in chunk_result.entities:
                    all_entities.append(entity)

            # Create a consolidated result for this document
            yield dspy.Prediction.from_completions({"entities": [all_entities]}, signature=self.prompt_signature)

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_template and _prompt_conclusion.

Note: different engines have different expectations as how a prompt should look like. E.g. outlines supports the Jinja 2 templating format for insertion of values and few-shot examples, whereas DSPy integrates these things in a different value in the workflow and hence expects the prompt not to include these things. Mind engine-specific expectations when creating a prompt template.

Returns:

Type	Description
str	Prompt template as string. None if not used by engine.

__init__(entities, task_id, prompt_instructions)

Initialize NERBridge.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID.	required
prompt_instructions	str | None	Custom prompt instructions. If None, default instructions are used.	required

Source code in sieves/tasks/predictive/ner/bridges.py

def __init__(
    self,
    entities: list[str],
    task_id: str,
    prompt_instructions: str | None,
):
    """Initialize NERBridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=False,
    )
    self._entities = entities

extract(docs)

Extract all values from doc instances that are to be injected into the prompts.

Overriding the default implementation to include the entity types in the extracted values.

Parameters:

Name	Type	Description	Default
docs	Iterable[Doc]	Docs to extract values from.	required

Returns:

Type	Description
Iterable[dict[str, Any]]	All values from doc instances that are to be injected into the prompts

Source code in sieves/tasks/predictive/ner/bridges.py

@override
def extract(self, docs: Iterable[Doc]) -> Iterable[dict[str, Any]]:
    """Extract all values from doc instances that are to be injected into the prompts.

    Overriding the default implementation to include the entity types in the extracted values.
    :param docs: Docs to extract values from.
    :return Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    return ({"text": doc.text if doc.text else None, "entity_types": self._entities} for doc in docs)

Entities

Bases: BaseModel

Collection of entities with associated text.

Source code in sieves/tasks/predictive/ner/bridges.py

class Entities(pydantic.BaseModel):
    """Collection of entities with associated text."""

    entities: list[Entity]
    text: str

Entity

Bases: BaseModel

Class for storing entity information.

Source code in sieves/tasks/predictive/ner/bridges.py

class Entity(pydantic.BaseModel):
    """Class for storing entity information."""

    text: str
    start: int
    end: int
    entity_type: str

    def __eq__(self, other: object) -> bool:
        """Compare two entities.

        :param other: Other entity to compare with.
        :return: True if entities are equal, False otherwise.
        """
        if not isinstance(other, Entity):
            return False
        # Two entities are equal if they have the same start, end, text and entity_type
        return (
            self.start == other.start
            and self.end == other.end
            and self.text == other.text
            and self.entity_type == other.entity_type
        )

    def __hash__(self) -> int:
        """Compute entity hash.

        :returns: Entity hash.
        """
        return hash((self.start, self.end, self.text, self.entity_type))

__eq__(other)

Compare two entities.

Parameters:

Name	Type	Description	Default
other	object	Other entity to compare with.	required

Returns:

Type	Description
bool	True if entities are equal, False otherwise.

Source code in sieves/tasks/predictive/ner/bridges.py

def __eq__(self, other: object) -> bool:
    """Compare two entities.

    :param other: Other entity to compare with.
    :return: True if entities are equal, False otherwise.
    """
    if not isinstance(other, Entity):
        return False
    # Two entities are equal if they have the same start, end, text and entity_type
    return (
        self.start == other.start
        and self.end == other.end
        and self.text == other.text
        and self.entity_type == other.entity_type
    )

__hash__()

Compute entity hash.

Returns:

Type	Description
int	Entity hash.

Source code in sieves/tasks/predictive/ner/bridges.py

def __hash__(self) -> int:
    """Compute entity hash.

    :returns: Entity hash.
    """
    return hash((self.start, self.end, self.text, self.entity_type))

GliXNER

Bases: NERBridge[list[str], Result, InferenceMode]

GliX bridge for NER.

Source code in sieves/tasks/predictive/ner/bridges.py

class GliXNER(NERBridge[list[str], glix_.Result, glix_.InferenceMode]):
    """GliX bridge for NER."""

    def __init__(
        self,
        entities: list[str],
        task_id: str,
        prompt_instructions: str | None,
    ):
        """Initialize GliXNER bridge.

        :param entities: List of entity types to extract.
        :param task_id: Task ID.
        :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
        """
        super().__init__(
            entities=entities,
            task_id=task_id,
            prompt_instructions=prompt_instructions,
        )

    @override
    @property
    def prompt_signature(self) -> list[str]:
        return self._entities

    @override
    @property
    def _default_prompt_instructions(self) -> str:
        return ""

    @override
    @property
    def _prompt_example_template(self) -> str | None:
        return None

    @override
    @property
    def _prompt_conclusion(self) -> str | None:
        return None

    @override
    @property
    def inference_mode(self) -> glix_.InferenceMode:
        return glix_.InferenceMode.ner

    @override
    def consolidate(
        self, results: Iterable[glix_.Result], docs_offsets: list[tuple[int, int]]
    ) -> Iterable[glix_.Result]:
        results = list(results)

        # Simply group results by document without trying to adjust positions
        # Position adjustment will happen in the integrate function
        for doc_offset in docs_offsets:
            doc_results = results[doc_offset[0] : doc_offset[1]]
            all_entities: list[dict[str, Any]] = []

            # Keep track of which chunk each entity came from
            for chunk_idx, chunk_result in enumerate(doc_results):
                # Process entities in this chunk
                for entity in chunk_result:
                    if isinstance(entity, dict):
                        # Add chunk index to the entity for reference in integrate
                        entity_copy = entity.copy()
                        entity_copy["chunk_idx"] = chunk_idx
                        all_entities.append(entity_copy)

            # Yield results for this document (flattened list of entities)
            yield all_entities

    @override
    def integrate(self, results: Iterable[glix_.Result], docs: Iterable[Doc]) -> Iterable[Doc]:
        docs_list = list(docs)
        results_list = list(results)

        # Process each document
        for doc, result in zip(docs_list, results_list):
            entities_list: list[Entity] = []
            doc_text = doc.text if doc.text is not None else ""

            # Get chunk information from the document
            chunk_offsets: list[int] = []
            if hasattr(doc, "chunks") and doc.chunks:
                # Calculate beginning position of each chunk in the original text
                current_offset = 0
                for chunk in doc.chunks:
                    chunk_offsets.append(current_offset)
                    current_offset += len(chunk) + 1

            # Process entities in this document
            if result:
                for entity_dict in result:
                    if not isinstance(entity_dict, dict):
                        continue

                    try:
                        entity_text = str(entity_dict.get("text", ""))
                        entity_start = int(entity_dict.get("start", 0))
                        entity_end = int(entity_dict.get("end", 0))
                        entity_type = str(entity_dict.get("label", ""))

                        # Get the chunk index (added in consolidate)
                        chunk_idx = int(entity_dict.get("chunk_idx", 0))

                        # Add chunk offset to entity positions
                        adjusted_start = entity_start
                        adjusted_end = entity_end

                        if chunk_offsets and chunk_idx < len(chunk_offsets):
                            # Adjust positions based on chunk offset
                            adjusted_start += chunk_offsets[chunk_idx]
                            adjusted_end += chunk_offsets[chunk_idx]

                        entities_list.append(
                            Entity(
                                text=entity_text,
                                start=adjusted_start,
                                end=adjusted_end,
                                entity_type=entity_type,
                            )
                        )
                    except (ValueError, TypeError) as e:
                        print(f"Error processing entity: {e}")
                        continue

            # Create the final entities object and store in document results
            entities_obj = Entities(text=doc_text, entities=entities_list)
            doc.results[self._task_id] = entities_obj

        return docs_list

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_template and _prompt_conclusion.

Note: different engines have different expectations as how a prompt should look like. E.g. outlines supports the Jinja 2 templating format for insertion of values and few-shot examples, whereas DSPy integrates these things in a different value in the workflow and hence expects the prompt not to include these things. Mind engine-specific expectations when creating a prompt template.

Returns:

Type	Description
str	Prompt template as string. None if not used by engine.

__init__(entities, task_id, prompt_instructions)

Initialize GliXNER bridge.

Parameters:

Name	Type	Description	Default
entities	list[str]	List of entity types to extract.	required
task_id	str	Task ID.	required
prompt_instructions	str | None	Custom prompt instructions. If None, default instructions are used.	required

Source code in sieves/tasks/predictive/ner/bridges.py

def __init__(
    self,
    entities: list[str],
    task_id: str,
    prompt_instructions: str | None,
):
    """Initialize GliXNER bridge.

    :param entities: List of entity types to extract.
    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    """
    super().__init__(
        entities=entities,
        task_id=task_id,
        prompt_instructions=prompt_instructions,
    )

extract(docs)

Extract all values from doc instances that are to be injected into the prompts.

Overriding the default implementation to include the entity types in the extracted values.

Parameters:

Name	Type	Description	Default
docs	Iterable[Doc]	Docs to extract values from.	required

Returns:

Type	Description
Iterable[dict[str, Any]]	All values from doc instances that are to be injected into the prompts

Source code in sieves/tasks/predictive/ner/bridges.py

@override
def extract(self, docs: Iterable[Doc]) -> Iterable[dict[str, Any]]:
    """Extract all values from doc instances that are to be injected into the prompts.

    Overriding the default implementation to include the entity types in the extracted values.
    :param docs: Docs to extract values from.
    :return Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    return ({"text": doc.text if doc.text else None, "entity_types": self._entities} for doc in docs)

LangChainNER

Bases: PydanticBasedNER[InferenceMode]

LangChain bridge for NER.

Source code in sieves/tasks/predictive/ner/bridges.py

class LangChainNER(PydanticBasedNER[langchain_.InferenceMode]):
    """LangChain bridge for NER."""

    @override
    @property
    def inference_mode(self) -> langchain_.InferenceMode:
        return langchain_.InferenceMode.structured

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_template and _prompt_conclusion.

Note: different engines have different expectations as how a prompt should look like. E.g. outlines supports the Jinja 2 templating format for insertion of values and few-shot examples, whereas DSPy integrates these things in a different value in the workflow and hence expects the prompt not to include these things. Mind engine-specific expectations when creating a prompt template.

Returns:

Type	Description
str	Prompt template as string. None if not used by engine.

__init__(entities, task_id, prompt_instructions)

Initialize NERBridge.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID.	required
prompt_instructions	str | None	Custom prompt instructions. If None, default instructions are used.	required

Source code in sieves/tasks/predictive/ner/bridges.py

def __init__(
    self,
    entities: list[str],
    task_id: str,
    prompt_instructions: str | None,
):
    """Initialize NERBridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=False,
    )
    self._entities = entities

extract(docs)

Extract all values from doc instances that are to be injected into the prompts.

Overriding the default implementation to include the entity types in the extracted values.

Parameters:

Name	Type	Description	Default
docs	Iterable[Doc]	Docs to extract values from.	required

Returns:

Type	Description
Iterable[dict[str, Any]]	All values from doc instances that are to be injected into the prompts

Source code in sieves/tasks/predictive/ner/bridges.py

@override
def extract(self, docs: Iterable[Doc]) -> Iterable[dict[str, Any]]:
    """Extract all values from doc instances that are to be injected into the prompts.

    Overriding the default implementation to include the entity types in the extracted values.
    :param docs: Docs to extract values from.
    :return Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    return ({"text": doc.text if doc.text else None, "entity_types": self._entities} for doc in docs)

NERBridge

Bases: Bridge[_BridgePromptSignature, _BridgeResult, EngineInferenceMode], ABC

Abstract base class for NER bridges.

Source code in sieves/tasks/predictive/ner/bridges.py

class NERBridge(Bridge[_BridgePromptSignature, _BridgeResult, EngineInferenceMode], abc.ABC):
    """Abstract base class for NER bridges."""

    def __init__(
        self,
        entities: list[str],
        task_id: str,
        prompt_instructions: str | None,
    ):
        """Initialize NERBridge.

        :param task_id: Task ID.
        :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
        """
        super().__init__(
            task_id=task_id,
            prompt_instructions=prompt_instructions,
            overwrite=False,
        )
        self._entities = entities

    @override
    def extract(self, docs: Iterable[Doc]) -> Iterable[dict[str, Any]]:
        """Extract all values from doc instances that are to be injected into the prompts.

        Overriding the default implementation to include the entity types in the extracted values.
        :param docs: Docs to extract values from.
        :return Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
        """
        return ({"text": doc.text if doc.text else None, "entity_types": self._entities} for doc in docs)

    @staticmethod
    def _find_entity_positions(
        doc_text: str,
        result: _BridgeResult,
    ) -> list[Entity]:
        """Find all positions of an entity in a document.

        :param doc_text: The text of the document.
        :param result: The result of the model.
        :return: The list of entities with start/end indices.
        """
        doc_text_lower = doc_text.lower()
        # Create a new result with the same structure as the original
        new_entities: list[Entity] = []

        # Track entities by position to avoid duplicates
        entities_by_position: dict[tuple[int, int], Entity] = {}
        context_list: list[str] = []

        entities_list = getattr(result, "entities", [])
        for entity_with_context in entities_list:
            # Skip if there is no entity
            if not entity_with_context:
                continue

            # Get the entity and context texts from the model
            entity_text = getattr(entity_with_context, "text", "")
            context = getattr(entity_with_context, "context", "")
            entity_type = getattr(entity_with_context, "entity_type", "")

            if not entity_text:
                continue

            entity_text_lower = entity_text.lower()
            context_lower = context.lower() if context else ""
            # Create a list of the unique contexts
            # Avoid adding duplicates as entities witht he same context would be captured twice
            if context_lower not in context_list:
                context_list.append(context_lower)
            else:
                continue
            # Find all occurrences of the context in the document using regex
            context_positions = re.finditer(re.escape(context_lower), doc_text_lower)

            # For each context position that was found (usually is just one), find the entity within that context
            for match in context_positions:
                context_start = match.start()
                entity_start_in_context = context_lower.find(entity_text_lower)

                if entity_start_in_context >= 0:
                    start = context_start + entity_start_in_context
                    end = start + len(entity_text)

                    # Create a new entity with start/end indices
                    new_entity = Entity(
                        text=doc_text[start:end],
                        start=start,
                        end=end,
                        entity_type=entity_type,
                    )

                    # Only add if this exact position hasn't been filled yet
                    position_key = (start, end)
                    if position_key not in entities_by_position:
                        entities_by_position[position_key] = new_entity
                        new_entities.append(new_entity)

        return sorted(new_entities, key=lambda x: x.start)

    @override
    def integrate(self, results: Iterable[_BridgeResult], docs: Iterable[Doc]) -> Iterable[Doc]:
        docs_list = list(docs)
        results_list = list(results)

        for doc, result in zip(docs_list, results_list):
            # Get the original text from the document
            doc_text = doc.text or ""
            if hasattr(result, "entities"):
                # Process entities from result if available
                entities_with_position = self._find_entity_positions(doc_text, result)
                # Create a new result with the updated entities
                new_result = Entities(text=doc_text, entities=entities_with_position)
                doc.results[self._task_id] = new_result
            else:
                # Default empty result
                doc.results[self._task_id] = Entities(text=doc_text, entities=[])

        return docs_list

inference_mode abstractmethod property

Return inference mode.

Returns:

Type	Description
EngineInferenceMode	Inference mode.

prompt_signature abstractmethod property

Create output signature.

E.g.: Signature in DSPy, Pydantic objects in outlines, JSON schema in jsonformers. This is engine-specific.

Returns:

Type	Description
type[TaskPromptSignature] | TaskPromptSignature	Output signature object. This can be an instance (e.g. a regex string) or a class (e.g. a Pydantic class).

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_template and _prompt_conclusion.

Note: different engines have different expectations as how a prompt should look like. E.g. outlines supports the Jinja 2 templating format for insertion of values and few-shot examples, whereas DSPy integrates these things in a different value in the workflow and hence expects the prompt not to include these things. Mind engine-specific expectations when creating a prompt template.

Returns:

Type	Description
str	Prompt template as string. None if not used by engine.

__init__(entities, task_id, prompt_instructions)

Initialize NERBridge.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID.	required
prompt_instructions	str | None	Custom prompt instructions. If None, default instructions are used.	required

Source code in sieves/tasks/predictive/ner/bridges.py

def __init__(
    self,
    entities: list[str],
    task_id: str,
    prompt_instructions: str | None,
):
    """Initialize NERBridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=False,
    )
    self._entities = entities

consolidate(results, docs_offsets) abstractmethod

Consolidate results for document chunks into document results.

Parameters:

Name	Type	Description	Default
results	Iterable[TaskResult]	Results per document chunk.	required
docs_offsets	list[tuple[int, int]]	Chunk offsets per document. Chunks per document can be obtained with results[docs_chunk_offsets[i][0]:docs_chunk_offsets[i][1]].	required

Returns:

Type	Description
Iterable[TaskResult]	Results per document.

Source code in sieves/tasks/predictive/bridges.py

@abc.abstractmethod
def consolidate(self, results: Iterable[TaskResult], docs_offsets: list[tuple[int, int]]) -> Iterable[TaskResult]:
    """Consolidate results for document chunks into document results.

    :param results: Results per document chunk.
    :param docs_offsets: Chunk offsets per document. Chunks per document can be obtained with
        `results[docs_chunk_offsets[i][0]:docs_chunk_offsets[i][1]]`.
    :return Iterable[_TaskResult]: Results per document.
    """

extract(docs)

Extract all values from doc instances that are to be injected into the prompts.

Overriding the default implementation to include the entity types in the extracted values.

Parameters:

Name	Type	Description	Default
docs	Iterable[Doc]	Docs to extract values from.	required

Returns:

Type	Description
Iterable[dict[str, Any]]	All values from doc instances that are to be injected into the prompts

Source code in sieves/tasks/predictive/ner/bridges.py

@override
def extract(self, docs: Iterable[Doc]) -> Iterable[dict[str, Any]]:
    """Extract all values from doc instances that are to be injected into the prompts.

    Overriding the default implementation to include the entity types in the extracted values.
    :param docs: Docs to extract values from.
    :return Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    return ({"text": doc.text if doc.text else None, "entity_types": self._entities} for doc in docs)

OutlinesNER

Bases: PydanticBasedNER[InferenceMode]

Outlines bridge for NER.

Source code in sieves/tasks/predictive/ner/bridges.py

class OutlinesNER(PydanticBasedNER[outlines_.InferenceMode]):
    """Outlines bridge for NER."""

    @override
    @property
    def inference_mode(self) -> outlines_.InferenceMode:
        return outlines_.InferenceMode.json

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_template and _prompt_conclusion.

Note: different engines have different expectations as how a prompt should look like. E.g. outlines supports the Jinja 2 templating format for insertion of values and few-shot examples, whereas DSPy integrates these things in a different value in the workflow and hence expects the prompt not to include these things. Mind engine-specific expectations when creating a prompt template.

Returns:

Type	Description
str	Prompt template as string. None if not used by engine.

__init__(entities, task_id, prompt_instructions)

Initialize NERBridge.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID.	required
prompt_instructions	str | None	Custom prompt instructions. If None, default instructions are used.	required

Source code in sieves/tasks/predictive/ner/bridges.py

def __init__(
    self,
    entities: list[str],
    task_id: str,
    prompt_instructions: str | None,
):
    """Initialize NERBridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=False,
    )
    self._entities = entities

extract(docs)

Extract all values from doc instances that are to be injected into the prompts.

Overriding the default implementation to include the entity types in the extracted values.

Parameters:

Name	Type	Description	Default
docs	Iterable[Doc]	Docs to extract values from.	required

Returns:

Type	Description
Iterable[dict[str, Any]]	All values from doc instances that are to be injected into the prompts

Source code in sieves/tasks/predictive/ner/bridges.py

@override
def extract(self, docs: Iterable[Doc]) -> Iterable[dict[str, Any]]:
    """Extract all values from doc instances that are to be injected into the prompts.

    Overriding the default implementation to include the entity types in the extracted values.
    :param docs: Docs to extract values from.
    :return Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    return ({"text": doc.text if doc.text else None, "entity_types": self._entities} for doc in docs)

PydanticBasedNER

Bases: NERBridge[BaseModel, BaseModel, EngineInferenceMode], ABC

Base class for Pydantic-based NER bridges.

Source code in sieves/tasks/predictive/ner/bridges.py

class PydanticBasedNER(NERBridge[pydantic.BaseModel, pydantic.BaseModel, EngineInferenceMode], abc.ABC):
    """Base class for Pydantic-based NER bridges."""

    @override
    @property
    def _default_prompt_instructions(self) -> str:
        return """
        Your goal is to extract named entities from the text. Only extract entities of the specified types:
        {{ entity_types }}.

        For each entity:
        - Extract the exact text of the entity
        - Include a SHORT context string that contains ONLY the entity and AT MOST 3 words before and 3 words after it.
          DO NOT include the entire text as context. DO NOT include words that are not present in the original text
          as introductory words (Eg. 'Text:' before context string).
        - Specify which type of entity it is (must be one of the provided entity types)

        IMPORTANT:
        - If the same entity appears multiple times in the text, extract each occurrence separately with its own context
        """

    @override
    @property
    def _prompt_example_template(self) -> str | None:
        return """
        {% if examples|length > 0 -%}
            <examples>
            {%- for example in examples %}
                <example>
                    <text>{{ example.text }}</text>
                    <entity_types>{{ entity_types }}</entity_types>
                    <entities>
                        {%- for entity in example.entities %}
                        <entity>
                            <text>{{ entity.text }}</text>
                            <context>{{ entity.context }}</context>
                            <entity_type>{{ entity.entity_type }}</entity_type>
                        </entity>
                        {%- endfor %}
                    </entities>
                </example>
            {% endfor -%}
            </examples>
        {% endif %}
        """

    @override
    @property
    def _prompt_conclusion(self) -> str | None:
        return """
        ===========

        <text>{{ text }}</text>
        <entity_types>{{ entity_types }}</entity_types>
        <entities>
        """

    @override
    @cached_property
    def prompt_signature(self) -> type[pydantic.BaseModel]:
        entity_types = self._entities
        LiteralType = Literal[*entity_types]  # type: ignore[valid-type]

        class EntityWithContext(pydantic.BaseModel):
            text: str
            context: str
            entity_type: LiteralType

        class Prediction(pydantic.BaseModel):
            """NER prediction."""

            entities: list[EntityWithContext] = []

        return Prediction

    @override
    def consolidate(
        self, results: Iterable[pydantic.BaseModel], docs_offsets: list[tuple[int, int]]
    ) -> Iterable[pydantic.BaseModel]:
        results = list(results)

        # Process each document (which may consist of multiple chunks)
        for doc_offset in docs_offsets:
            doc_results = results[doc_offset[0] : doc_offset[1]]

            # Combine all entities from all chunks
            all_entities: list[dict[str, Any]] = []

            # Process each chunk for this document
            for chunk_result in doc_results:
                if not hasattr(chunk_result, "entities") or not chunk_result.entities:
                    continue

                # Process entities in this chunk
                for entity in chunk_result.entities:
                    # We just need to combine all entities from all chunks
                    all_entities.append(entity)

            # Create a consolidated result for this document - instantiate the class with entities
            yield self.prompt_signature(entities=all_entities)

inference_mode abstractmethod property

Return inference mode.

Returns:

Type	Description
EngineInferenceMode	Inference mode.

prompt_template property

Return prompt template.

Chains _prompt_instructions, _prompt_example_template and _prompt_conclusion.

Note: different engines have different expectations as how a prompt should look like. E.g. outlines supports the Jinja 2 templating format for insertion of values and few-shot examples, whereas DSPy integrates these things in a different value in the workflow and hence expects the prompt not to include these things. Mind engine-specific expectations when creating a prompt template.

Returns:

Type	Description
str	Prompt template as string. None if not used by engine.

__init__(entities, task_id, prompt_instructions)

Initialize NERBridge.

Parameters:

Name	Type	Description	Default
task_id	str	Task ID.	required
prompt_instructions	str | None	Custom prompt instructions. If None, default instructions are used.	required

Source code in sieves/tasks/predictive/ner/bridges.py

def __init__(
    self,
    entities: list[str],
    task_id: str,
    prompt_instructions: str | None,
):
    """Initialize NERBridge.

    :param task_id: Task ID.
    :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
    """
    super().__init__(
        task_id=task_id,
        prompt_instructions=prompt_instructions,
        overwrite=False,
    )
    self._entities = entities

extract(docs)

Extract all values from doc instances that are to be injected into the prompts.

Overriding the default implementation to include the entity types in the extracted values.

Parameters:

Name	Type	Description	Default
docs	Iterable[Doc]	Docs to extract values from.	required

Returns:

Type	Description
Iterable[dict[str, Any]]	All values from doc instances that are to be injected into the prompts

Source code in sieves/tasks/predictive/ner/bridges.py

@override
def extract(self, docs: Iterable[Doc]) -> Iterable[dict[str, Any]]:
    """Extract all values from doc instances that are to be injected into the prompts.

    Overriding the default implementation to include the entity types in the extracted values.
    :param docs: Docs to extract values from.
    :return Iterable[dict[str, Any]]: All values from doc instances that are to be injected into the prompts
    """
    return ({"text": doc.text if doc.text else None, "entity_types": self._entities} for doc in docs)