Merge branch 'dev'

Author: Aethor

docs/reference.rst

@@ -64,6 +64,13 @@ NLTKNamedEntityRecognizer
    :members:
 
 
+Retrieval
+---------
+
+.. automodule:: renard.pipeline.ner.retrieval
+   :members:
+
+
 Coreference Resolution
 ======================
 

pyproject.toml

@@ -1,31 +1,31 @@
 [project]
 name = "renard-pipeline"
-version = "0.7.1"
+version = "0.7.2"
 description = "Relationships Extraction from NARrative Documents"
 authors = [
     {name = "Arthur Amalvy", email = "arthur.amalvy@univ-avignon.fr"},
 ]
 license = { text = "GPL-3.0-only" }
 readme = "README.md"
-requires-python = ">=3.9,<3.13"
+requires-python = ">=3.10,<3.14"
 dependencies = [
-    "torch>=2.7.0",
-    "transformers>=4.57.1",
-    "nltk>=3.9.1",
+    "torch>=2.8.0",
+    "transformers>=5.0.0",
+    "nltk>=3.10.3",
     "tqdm>=4.67.1",
     "networkx>=3.2",
     "more-itertools>=10.7",
     "nameparser>=1.1.3",
     "matplotlib>=3.9",
     "pytest>=8.4.1",
-    "tibert>=0.5.2",
+    "tibert>=0.5.3",
     "grimbert>=0.1.5",
     "datasets>=4.0.0",
     "rank-bm25>=0.2.2",
     "accelerate>=1.10.1",
     "scikit-learn>=1.6.1",
     "tiktoken>=0.12.0",
-    "protobuf>=6.33.2",
+    "protobuf>=6.33.5",
 ]
 
 [build-system]
@@ -50,44 +50,67 @@ dev = [
 
 [project.optional-dependencies]
 ui = [
-    "gradio>=4.44.1",
+    "gradio>=6.0.0",
     "pyvis>=0.3.2",
 ]
 cpu = [
-  "torch>=2.7.1",
+  "torch>=2.8.0",
 ]
 cuda128 = [
-  "torch>=2.7.1",
+  "torch>=2.8.0",
+]
+cuda129 = [
+  "torch>=2.8.0",
 ]
 rocm63 = [
-  "torch>=2.7.1",
-  "pytorch-triton-rocm>=3.1.0", 
+  "torch>=2.8.0",
+  "triton-rocm",
+  "pytorch-triton-rocm>=3.5.0 ; sys_platform == 'linux'", 
 ]
 rocm64 = [
-  "torch>=2.7.1",
-  "pytorch-triton-rocm>=3.1.0", 
+  "torch>=2.8.0",
+  "triton-rocm",
+  "pytorch-triton-rocm>=3.5.0 ; sys_platform == 'linux'", 
+]
+rocm71 = [
+  "torch>=2.8.0",
+  "triton-rocm",
+  "pytorch-triton-rocm>=3.1.0 ; sys_platform == 'linux'", 
+]
+dspy = [
+    "dspy>=2.6.27",
 ]
 
 [tool.uv]
 conflicts = [
   [
     { extra = "cpu" },
     { extra = "cuda128" },
+    { extra = "cuda129" },
     { extra = "rocm63" },
     { extra = "rocm64" },
+    { extra = "rocm71" },
   ],
 ]
 
 [tool.uv.sources]
 torch = [
   { index = "pytorch-cpu", extra = "cpu" },
   { index = "pytorch-cuda128", extra = "cuda128" },
-  { index = "pytorch-rocm63", extra = "rocm63" },
-  { index = "pytorch-rocm64", extra = "rocm64" },
+  { index = "pytorch-cuda129", extra = "cuda129" },
+  { index = "pytorch-rocm63", extra = "rocm63", marker = "sys_platform == 'linux'" },
+  { index = "pytorch-rocm64", extra = "rocm64", marker = "sys_platform == 'linux'" },
+  { index = "pytorch-rocm71", extra = "rocm71", marker = "sys_platform == 'linux'" },
 ]
 pytorch-triton-rocm = [
-  { index = "pytorch-rocm63", extra = "rocm63" },
-  { index = "pytorch-rocm64", extra = "rocm64" },
+  { index = "pytorch-rocm63", extra = "rocm63", marker = "sys_platform == 'linux'" },
+  { index = "pytorch-rocm64", extra = "rocm64", marker = "sys_platform == 'linux'" },
+  { index = "pytorch-rocm71", extra = "rocm71", marker = "sys_platform == 'linux'" },
+]
+triton-rocm = [
+  { index = "pytorch-rocm63", extra = "rocm63", marker = "sys_platform == 'linux'" },
+  { index = "pytorch-rocm64", extra = "rocm64", marker = "sys_platform == 'linux'" },
+  { index = "pytorch-rocm71", extra = "rocm71", marker = "sys_platform == 'linux'" },
 ]
 
 [[tool.uv.index]]
@@ -100,6 +123,11 @@ name = "pytorch-cuda128"
 url = "https://download.pytorch.org/whl/cu128"
 explicit = true
 
+[[tool.uv.index]]
+name = "pytorch-cuda129"
+url = "https://download.pytorch.org/whl/cu129"
+explicit = true
+
 [[tool.uv.index]]
 name = "pytorch-rocm63"
 url = "https://download.pytorch.org/whl/rocm6.3"
@@ -108,4 +136,9 @@ explicit = true
 [[tool.uv.index]]
 name = "pytorch-rocm64"
 url = "https://download.pytorch.org/whl/rocm6.4"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-rocm71"
+url = "https://download.pytorch.org/whl/rocm7.1"
 explicit = true

renard/ner_utils.py

@@ -6,7 +6,7 @@
 from more_itertools import flatten
 import torch
 from torch.utils.data import Dataset
-from datasets import Dataset as HGDataset
+from datasets import Dataset as HFDataset, DatasetDict as HFDatasetDict
 from datasets import Sequence, ClassLabel
 from transformers import (
     AutoModelForTokenClassification,
@@ -92,7 +92,7 @@ def __init__(
             assert all(
                 [len(cm) == len(elt) for elt, cm in zip(self.elements, context_mask)]
             )
-        self._context_mask = context_mask or [[1] * len(elt) for elt in self.elements]
+        self._context_mask = context_mask or [[0] * len(elt) for elt in self.elements]
 
         self.tokenizer = tokenizer
 
@@ -213,8 +213,10 @@ def load_conll2002_bio(
     tags = []
     for line in raw_data.split("\n"):
         line = line.strip("\n")
-        if re.fullmatch(r"\s*", line) or (
-            not max_sent_len is None and len(sent_tokens) >= max_sent_len
+        if (
+            re.fullmatch(r"\s*", line)  # ignore empty lines
+            or re.fullmatch(r"# [^:]+: .*", line)  # ignore Novelties style metadata
+            or (not max_sent_len is None and len(sent_tokens) >= max_sent_len)
         ):
             if len(sent_tokens) == 0:
                 continue
@@ -224,29 +226,36 @@ def load_conll2002_bio(
         token, tag = line.split(separator)
         sent_tokens.append(token)
         tags.append(tag_conversion_map.get(tag, tag))
+    if len(sent_tokens) != 0:
+        sents.append(sent_tokens)
 
     tokens = list(flatten(sents))
     entities = ner_entities(tokens, tags)
 
     return sents, list(flatten(sents)), entities
 
 
-def hgdataset_from_conll2002(
+def hfdataset_from_conll2002(
     path: str,
     tag_conversion_map: Optional[Dict[str, str]] = None,
     separator: str = "\t",
     max_sent_len: Optional[int] = None,
+    labels: Optional[list[str]] = None,
     **kwargs,
-) -> HGDataset:
+) -> HFDataset:
     """Load a CoNLL-2002 file as a Huggingface Dataset.
 
     :param path: passed to :func:`.load_conll2002_bio`
     :param tag_conversion_map: passed to :func:`load_conll2002_bio`
     :param separator: passed to :func:`load_conll2002_bio`
     :param max_sent_len: passed to :func:`load_conll2002_bio`
+    :param labels: the list of all possible labels.  If ``None``, will
+        automatically be assigned to the sorted list of possible tags
+        found in the input file.
     :param kwargs: additional kwargs for :func:`open`
 
-    :return: a :class:`datasets.Dataset` with features 'tokens' and 'labels'.
+    :return: a :class:`datasets.Dataset` with features 'tokens' and
+             'labels'.
     """
     sentences, tokens, entities = load_conll2002_bio(
         path, tag_conversion_map, separator, max_sent_len, **kwargs
@@ -268,13 +277,21 @@ def hgdataset_from_conll2002(
         for sent_start, sent_end in zip(sent_starts, sent_ends)
     ]
 
-    dataset = HGDataset.from_dict({"tokens": sentences, "labels": sent_tags})
-    dataset = dataset.cast_column(
-        "labels", Sequence(ClassLabel(names=sorted(set(tags))))
-    )
+    dataset = HFDataset.from_dict({"tokens": sentences, "labels": sent_tags})
+    if labels is None:
+        labels = sorted(set(tags))
+    dataset = dataset.cast_column("labels", Sequence(ClassLabel(names=labels)))
     return dataset
 
 
+def hgdataset_from_conll2002(**kwargs) -> HFDataset:
+    """
+    Deprecated function that only exists for retrocompatibility, you
+    should call :func:`.hfdataset_from_conll2002` instead.
+    """
+    return hfdataset_from_conll2002(**kwargs)
+
+
 def _tokenize_and_align_labels(
     examples, tokenizer: PreTrainedTokenizerFast, label_all_tokens: bool = True
 ):
@@ -315,37 +332,49 @@ def _tokenize_and_align_labels(
 
 
 def train_ner_model(
-    hg_id: str,
-    dataset: HGDataset,
+    hf_id: str,
+    dataset: Union[HFDataset, HFDatasetDict],
     targs: TrainingArguments,
+    train_split: str = "train",
+    valid_split: str = "valid",
+    trainer_class: type[Trainer] = Trainer,
 ) -> PreTrainedModel:
+    """Train a NER model on the given dataset.
+
+    :param hf_id: huggingface ID of the model to train
+    :param dataset: huggingface dataset on which to train.  The
+        'labels' column is assumed to contain NER labels.
+    :param TrainingArguments: training arguments for the huggingface
+        trainer.
+    :param train_split: split of the dataset used for train.
+    :param valid_split: split of the dataset used for validation.
+    :param trainer_class: trainer class to use.  Can be used to
+        override the default huggingface trainer.
+    """
     from transformers import DataCollatorForTokenClassification
 
     # BERT tokenizer splits tokens into subtokens. The
     # tokenize_and_align_labels function correctly aligns labels and
     # subtokens.
-    tokenizer = AutoTokenizer.from_pretrained(hg_id)
+    tokenizer = AutoTokenizer.from_pretrained(hf_id)
     dataset = dataset.map(
         ft.partial(_tokenize_and_align_labels, tokenizer=tokenizer), batched=True
     )
-    dataset = dataset.train_test_split(test_size=0.1)
 
-    label_lst = dataset["train"].features["labels"].feature.names
+    label_lst = dataset[train_split].features["labels"].feature.names
     model = AutoModelForTokenClassification.from_pretrained(
-        hg_id,
+        hf_id,
         num_labels=len(label_lst),
         id2label={i: label for i, label in enumerate(label_lst)},
         label2id={label: i for i, label in enumerate(label_lst)},
     )
 
-    trainer = Trainer(
+    trainer = trainer_class(
         model,
         targs,
-        train_dataset=dataset["train"],
-        eval_dataset=dataset["test"],
-        # data_collator=DataCollatorForTokenClassificationWithBatchEncoding(tokenizer),
+        train_dataset=dataset[train_split],
+        eval_dataset=dataset[valid_split],
         data_collator=DataCollatorForTokenClassification(tokenizer),
-        tokenizer=tokenizer,
     )
     trainer.train()
 

renard/pipeline/core.py

@@ -82,7 +82,7 @@ def __init__(self):
 
     def _pipeline_init_(
         self, lang: str, progress_reporter: ProgressReporter, **kwargs
-    ) -> Optional[Dict[Pipeline.PipelineParameter, Any]]:
+    ) -> Optional[Dict[str, Any]]:
         """Set the step configuration that is common to the whole
         pipeline.
 
@@ -576,16 +576,13 @@ def update(slider_value):
 class Pipeline:
     """A flexible NLP pipeline"""
 
-    #: all the possible parameters of the whole pipeline, that are
-    #: shared between steps
-    PipelineParameter = Literal["lang", "progress_reporter", "character_ner_tag"]
-
     def __init__(
         self,
         steps: List[PipelineStep],
         lang: str = "eng",
         progress_report: Optional[Literal["tqdm"]] = "tqdm",
         warn: bool = True,
+        **step_additional_params,
     ) -> None:
         """
         :param steps: a ``tuple`` of :class:``PipelineStep``, that
@@ -595,16 +592,28 @@ def __init__(
             progress.
         :param lang: ISO 639-3 language code
         :param warn:
+
+        :param step_additional_params: additional parameters passed to
+            :meth:`._pipeline_init_` when
+            :meth:`_pipeline_init_steps_` is called.  The following
+            values are currently used:
+
+                - ``'character_ner_tag'``: the NER tag corresponding
+                  to characters (default: ``PER``)
         """
         self.steps = steps
 
         self.progress_report: Optional[Literal["tqdm"]] = progress_report
         self.progress_reporter = get_progress_reporter(progress_report)
 
         self.lang = lang
-        self.character_ner_tag = "PER"
         self.warn = warn
 
+        self.step_additional_params = step_additional_params
+        self.step_additional_params["character_ner_tag"] = (
+            self.step_additional_params.get("character_ner_tag", "PER")
+        )
+
     def _pipeline_init_steps_(self, ignored_steps: Optional[List[str]] = None):
         """Initialise steps with global pipeline parameters.
 
@@ -613,16 +622,18 @@ def _pipeline_init_steps_(self, ignored_steps: Optional[List[str]] = None):
         """
         steps_progress_reporter = self.progress_reporter.get_subreporter()
         steps = self._non_ignored_steps(ignored_steps)
-        pipeline_params = {
-            "progress_reporter": steps_progress_reporter,
-            "character_ner_tag": self.character_ner_tag,
-        }
+        pipeline_params = self.step_additional_params.copy()
         for step in steps:
-            step_additional_params = step._pipeline_init_(self.lang, **pipeline_params)
+            step_additional_params = step._pipeline_init_(
+                self.lang, progress_reporter=steps_progress_reporter, **pipeline_params
+            )
             if not step_additional_params is None:
                 for key, value in step_additional_params.items():
                     setattr(self, key, value)
-                    pipeline_params[key] = value
+                    # parameters set by the user have precedence over
+                    # step mandated parameters
+                    if not key in self.step_additional_params:
+                        pipeline_params[key] = value
 
     def _non_ignored_steps(
         self, ignored_steps: Optional[List[str]]