llm-annotator 0.3.2__tar.gz → 0.3.4__tar.gz

{llm_annotator-0.3.2 → llm_annotator-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llm-annotator
-Version: 0.3.2
+Version: 0.3.4
 Summary: An easy-to-extend LLM annotator for robust, resumable data annotation.
 Author-email: Bram Vanroy <2779410+BramVanroy@users.noreply.github.com>
 License-Expression: Apache-2.0

{llm_annotator-0.3.2 → llm_annotator-0.3.4}/examples/sentiment.py

@@ -35,6 +35,12 @@ Classification:"""
 
     def random_validitity(sample):
         return random.random() < 0.5
+    
+    def postprocess_fn(sample):
+        # Example postprocessing: strip whitespace from sentiment
+        if "sentiment" in sample and isinstance(sample["sentiment"], str):
+            sample["sentiment"] = sample["sentiment"].strip()
+        return sample
 
     with Annotator(model="Qwen/Qwen2.5-0.5B-Instruct", max_model_len=4096, verbose=True) as anno:
         ds = anno.annotate_dataset(
@@ -55,6 +61,7 @@ Classification:"""
             sort_by_length=True,  # Sort by prompt length for more efficient batching -- final dataset will be re-ordered to original
             validate_fn=random_validitity,
             num_retries_invalid=3,
+            postprocess_fn=postprocess_fn,
         )
     print(ds)
     shutil.rmtree("outputs/sentiment-imdb-qwen")

{llm_annotator-0.3.2 → llm_annotator-0.3.4}/src/llm_annotator/annotator.py

@@ -17,7 +17,7 @@ from vllm import LLM, RequestOutput, SamplingParams
 from vllm.distributed import destroy_distributed_environment, destroy_model_parallel
 from vllm.sampling_params import StructuredOutputsParams
 
-from llm_annotator.utils import ensure_returns_bool, remove_empty_jsonl_files, retry
+from llm_annotator.utils import ensure_returns_bool, ensure_returns_dict, remove_empty_jsonl_files, retry
 
 
 def destroy_model_on_error(func):
@@ -385,15 +385,30 @@ class Annotator:
         """Load and configure the tokenizer for the model.
 
         Sets up the tokenizer with appropriate padding settings and ensures
-        a pad token is available.
+        a pad token is available. Some models (like recent Mistral) do not have
+        a chat template defined in their tokenizer, so we attempt to load
+        a processor instead in that case.
         """
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model)
-
-        self.tokenizer.padding_side = "left"
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model)
+            if self.tokenizer.chat_template is None:
+                raise ValueError("No chat template defined in the tokenizer.")
+        except ValueError as tokenizer_error:
+            try:
+                self.tokenizer = self.AutoProcessor.from_pretrained(self.model)
+                if self.tokenizer.chat_template is None:
+                    raise ValueError("No chat template defined in the processor either.") from tokenizer_error
+            except Exception as exc:
+                raise ValueError(
+                    "Failed to load tokenizer or processor with chat template."
+                    " Are you correctly using an instruct model?"
+                ) from exc
+        else:
+            self.tokenizer.padding_side = "left"
 
-        if not self.tokenizer.pad_token_id:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-            self.tokenizer.pad_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.pad_token)
+            if not self.tokenizer.pad_token_id:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                self.tokenizer.pad_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.pad_token)
 
     def _load_pipeline(self) -> None:
         """Load and initialize the vLLM pipeline for inference.
@@ -489,6 +504,7 @@ class Annotator:
         sampling_params: SamplingParams,
         task_prefix: str = "",
         validate_fn: Callable | None = None,
+        postprocess_fn: Callable | None = None,
     ) -> list[dict[str, Any]]:
         """Process a batch of samples through the model.
 
@@ -506,23 +522,28 @@ class Annotator:
                 with keys as produced by `_process_output`: 1. the raw response ("{prefix}response"),
                 2. the finish reason ("{prefix}finish_reason"), 3. the number of tokens ("{prefix}num_tokens").
                 When an output schema was provided, also the parsed JSON fields and the "{prefix}valid_fields" key.
+            postprocess_fn: Optional function to postprocess each sample after annotation. Must return a modified
+                sample dictionary.
 
         Returns:
             List of processed output dictionaries for each sample in the batch.
         """
         output_schema = sampling_params.structured_outputs.json if sampling_params.structured_outputs else None
         outputs = self.pipe.generate(batch[f"{task_prefix}prompted"], sampling_params, use_tqdm=False)
-        results = [
-            self._process_output(output=outp, output_schema=output_schema, task_prefix=task_prefix) for outp in outputs
-        ]
 
-        if validate_fn:
-            for res in results:
+        results = []
+        for outp in outputs:
+            res = self._process_output(output=outp, output_schema=output_schema, task_prefix=task_prefix)
+            if postprocess_fn:
+                res = ensure_returns_dict(postprocess_fn, res)
+
+            if validate_fn:
                 if f"{task_prefix}valid_fields" in res and res[f"{task_prefix}valid_fields"] is False:
                     is_valid = False
                 else:
                     is_valid = ensure_returns_bool(validate_fn, res)
                 res[f"{task_prefix}valid"] = is_valid
+            results.append(res)
 
         if f"{task_prefix}valid_fields" in results[0]:
             n_invalid = sum([1 for res in results if not res[f"{task_prefix}valid_fields"]])
@@ -576,7 +597,9 @@ class Annotator:
         task_prefix: str = "",
         sort_by_length: bool = False,
         validate_fn: Callable | None = None,
-        num_retries_invalid: int = 0,
+        postprocess_fn: Callable | None = None,
+        num_retries_invalid: int = 5,
+        keep_idx_column: bool = False,
     ) -> Dataset:
         """Annotate an entire dataset using the configured model and prompt.
 
@@ -631,11 +654,16 @@ class Annotator:
                 with keys as produced by `_process_output`: 1. the raw response ("{prefix}response"),
                 2. the finish reason ("{prefix}finish_reason"), 3. the number of tokens ("{prefix}num_tokens").
                 When an output schema was provided, also the parsed JSON fields and the "{prefix}valid_fields" key.
+            postprocess_fn: Optional function to postprocess each sample after annotation. Must return a modified
+                sample dictionary.
             num_retries_invalid: Number of retries for samples that produce invalid outputs (when
                 a JSON schema is given and {prefix}valid_fields is False, or when a validate_fn is given
                 and it {prefix}valid is False).
+            keep_idx_column: Whether to keep the idx_column in the final dataset before uploading and returning.
+
+        Returns:
+            The concatenated dataset of all annotation results (JSON-invalid samples are NOT removed)
         """
-        # Verify shared_prompt_template
         if prompt_template_prefix:
             if prompt_template_prefix not in full_prompt_template:
                 raise ValueError(
@@ -763,7 +791,11 @@ class Annotator:
                 unit="batch",
             ):
                 results = self._process_batch(
-                    batch=batch, sampling_params=sampling_params, task_prefix=task_prefix, validate_fn=validate_fn
+                    batch=batch,
+                    sampling_params=sampling_params,
+                    task_prefix=task_prefix,
+                    validate_fn=validate_fn,
+                    postprocess_fn=postprocess_fn,
                 )
 
                 if num_retries_invalid > 0:
@@ -846,9 +878,13 @@ class Annotator:
             if new_hub_id and upload_every_n_samples > 0:
                 self.push_dir_to_hub(pdout, new_hub_id=new_hub_id)
 
-        return self._post_annotate(pdout=pdout, idx_column=idx_column, new_hub_id=new_hub_id)
+        return self._post_annotate(
+            pdout=pdout, idx_column=idx_column, new_hub_id=new_hub_id, keep_idx_column=keep_idx_column
+        )
 
-    def _post_annotate(self, *, pdout: Path, idx_column: str, new_hub_id: str | None = None) -> Dataset:
+    def _post_annotate(
+        self, *, pdout: Path, idx_column: str, new_hub_id: str | None = None, keep_idx_column: bool = False
+    ) -> Dataset:
         """Clean up after annotation is complete.
 
         Removes empty output files and performs any final cleanup operations.
@@ -856,6 +892,8 @@ class Annotator:
         Args:
             pdout: Output directory path to clean up.
             new_hub_id: Optional Hugging Face dataset ID for uploads.
+            idx_column: Column name used as unique identifier.
+            keep_idx_column: Whether to keep the idx_column in the final dataset before uploading and returning
 
         Returns:
             The concatenated dataset of all annotation results (JSON-invalid samples are NOT removed)
@@ -865,7 +903,9 @@ class Annotator:
             if pfin.stat().st_size > 0:
                 ds_parts.append(Dataset.from_json(str(pfin)))
 
-        ds: Dataset = concatenate_datasets(ds_parts).sort(idx_column).remove_columns([idx_column])
+        ds: Dataset = concatenate_datasets(ds_parts).sort(idx_column)
+        if not keep_idx_column:
+            ds = ds.remove_columns([idx_column])
 
         if new_hub_id:
             ds.push_to_hub(new_hub_id, private=True)

{llm_annotator-0.3.2 → llm_annotator-0.3.4}/src/llm_annotator/utils.py

@@ -155,3 +155,11 @@ def ensure_returns_bool(func, *args, **kwargs):
     if not isinstance(result, bool):
         raise TypeError(f"{func.__name__} should return a bool, got {type(result).__name__}")
     return result
+
+
+def ensure_returns_dict(func, *args, **kwargs):
+    """Ensure that the given function returns a dict value. If not, raise a TypeError."""
+    result = func(*args, **kwargs)
+    if not isinstance(result, dict):
+        raise TypeError(f"{func.__name__} should return a dict, got {type(result).__name__}")
+    return result