llm-annotator 0.2.4__tar.gz → 0.2.6__tar.gz

{llm_annotator-0.2.4 → llm_annotator-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llm-annotator
-Version: 0.2.4
+Version: 0.2.6
 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.2.4 → llm_annotator-0.2.6}/examples/sentiment.py

@@ -32,11 +32,11 @@ def main():
             dataset_split="test",
             new_hub_id=f"{hf_user}/sentiment-imdb",
             streaming=True,
-            max_num_samples=250,
+            max_num_samples=200,
             cache_input_dataset=False,  # `True` is generally useful, not for demo purposes
             prompt_template=prompt_template,
             output_schema=output_schema,
-            # Backup to HF every 10 samples.
+            # Backup to HF every 100 samples.
             # In practice, set to a higher value (e.g., 1000+)
             upload_every_n_samples=100,
         )

{llm_annotator-0.2.4 → llm_annotator-0.2.6}/src/llm_annotator/annotator.py

@@ -3,6 +3,7 @@ import json
 import shutil
 import string
 from dataclasses import dataclass, field
+from functools import wraps
 from math import ceil
 from os import PathLike
 from pathlib import Path
@@ -15,11 +16,29 @@ from tqdm import tqdm
 from transformers import AutoTokenizer, PreTrainedTokenizer
 from vllm import LLM, RequestOutput, SamplingParams
 from vllm.distributed import destroy_distributed_environment, destroy_model_parallel
-from vllm.sampling_params import GuidedDecodingParams
+from vllm.sampling_params import StructuredOutputsParams
 
 from llm_annotator.utils import remove_empty_jsonl_files, retry
 
 
+def destroy_model_on_error(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        try:
+            return func(self, *args, **kwargs)
+        except BaseException as e:
+            # Catch BaseException so we also clean on KeyboardInterrupt/SystemExit
+            try:
+                self.destroy_model()
+            except Exception as clean_err:
+                # Don't hide the original error; attach a note (Python 3.11+)
+                if hasattr(e, "__notes__"):
+                    e.__notes__.append(f"Cleanup failed: {clean_err!r}")
+            raise  # re-raise the original exception
+
+    return wrapper
+
+
 @dataclass(slots=True)
 class Annotator:
     """Sensible base class for LLM-based dataset annotation.
@@ -413,16 +432,21 @@ class Annotator:
         """
         if isinstance(self.pipe, LLM):
             destroy_model_parallel()
+            destroy_distributed_environment()
+
             try:
                 self.pipe.llm_engine.model_executor.shutdown()
                 del self.pipe.llm_engine.model_executor
             except Exception:
                 pass
+            try:
+                self.pipe.llm_engine.engine_core.shutdown()
+                del self.pipe.llm_engine.engine_core
+            except Exception:
+                pass
             del self.pipe.llm_engine
             del self.pipe
-            gc.collect()
             cuda.empty_cache()
-            destroy_distributed_environment()
             gc.collect()
 
             self.pipe = None
@@ -447,7 +471,7 @@ class Annotator:
         Returns:
             List of processed output dictionaries for each sample in the batch.
         """
-        output_schema = sampling_params.guided_decoding.json if sampling_params.guided_decoding else None
+        output_schema = sampling_params.structured_outputs.json if sampling_params.structured_outputs else None
         outputs = self.pipe.generate(batch[f"{prefix}prompted"], sampling_params, use_tqdm=False)
         results = [self._process_output(output=outp, output_schema=output_schema, prefix=prefix) for outp in outputs]
 
@@ -462,6 +486,7 @@ class Annotator:
 
         return results
 
+    @destroy_model_on_error
     def annotate_dataset(
         self,
         output_dir: str | Path,
@@ -632,11 +657,11 @@ class Annotator:
 
             sampling_params = sampling_params or {}
             if output_schema:
-                if "guided_decoding" in sampling_params:
+                if "structured_outputs" in sampling_params:
                     raise ValueError(
-                        "If 'output_schema' is provided, do not set 'guided_decoding' in sampling_params yourself"
+                        "If 'output_schema' is provided, do not set 'structured_outputs' in sampling_params yourself"
                     )
-                sampling_params["guided_decoding"] = GuidedDecodingParams(
+                sampling_params["structured_outputs"] = StructuredOutputsParams(
                     json=output_schema,
                     whitespace_pattern=whitespace_pattern,
                 )