llm-annotator 0.2.2__tar.gz → 0.2.3__tar.gz

llm_annotator-0.2.3/Makefile

@@ -0,0 +1,11 @@
+quality:
+	ruff check src/llm_annotator tests/ examples/
+	ruff format --check src/llm_annotator tests/ examples/
+
+style:
+	ruff check src/llm_annotator tests/ examples/ --fix
+	ruff format src/llm_annotator tests/ examples/
+
+setup:
+	uv sync --dev
+	pre-commit install --hook-type pre-push

{llm_annotator-0.2.2 → llm_annotator-0.2.3}/PKG-INFO

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

@@ -1,5 +1,3 @@
-
-import json
 from huggingface_hub import HfApi
 
 from llm_annotator import Annotator
@@ -11,7 +9,8 @@ def get_hf_username() -> str | None:
         return whoami["name"]
     else:
         raise ValueError("No Hugging Face username found. Please login using `hf auth login`.")
-    
+
+
 def main():
     hf_user = get_hf_username()
     prompt_template = """Analyze the sentiment of the following movie review and classify it as positive or negative.
@@ -22,24 +21,19 @@ def main():
 
     output_schema = {
         "type": "object",
-        "properties": {
-            "sentiment": {
-                "type": "string",
-                "enum": ["positive", "negative", "neutral"]
-            }
-        },
-        "required": ["sentiment"]
+        "properties": {"sentiment": {"type": "string", "enum": ["positive", "negative", "neutral"]}},
+        "required": ["sentiment"],
     }
 
     with Annotator(model_id="Qwen/Qwen2.5-0.5B-Instruct") as anno:
         ds = anno.annotate_dataset(
-            "stanfordnlp/imdb",
             output_dir="outputs/sentiment-imdb-qwen",
+            dataset_name="stanfordnlp/imdb",
             dataset_split="test",
             new_hub_id=f"{hf_user}/sentiment-imdb",
             streaming=True,
             max_num_samples=250,
-            cache_input_dataset=False,  # `True` is generally useful, not for demo purposes        
+            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.
@@ -48,5 +42,6 @@ def main():
         )
     print(ds)
 
+
 if __name__ == "__main__":
-    main()
+    main()

{llm_annotator-0.2.2 → llm_annotator-0.2.3}/src/llm_annotator/annotator.py

@@ -62,7 +62,7 @@ class Annotator:
         return self
 
     def __exit__(self, exc_type, exc, tb):
-        self.reset_model()
+        self.destroy_model()
 
     def _get_skip_idxs(
         self, *, pdout: Path, idx_column: str, dataset_split: str | None, dataset_config: str | None
@@ -406,29 +406,21 @@ class Annotator:
             **result,
         }
 
-    def reset_model(self) -> None:
+    def destroy_model(self) -> None:
         """Clean up model resources to free memory.
 
-        Destroys the distributed environment, clears GPU cache, and resets internal state..
+        Destroys the distributed environment, clears GPU cache, and resets internal state.
         """
-        destroy_model_parallel()
-        destroy_distributed_environment()
+        if isinstance(self.pipe, LLM):
+            destroy_model_parallel()
+            del self.pipe.llm_engine
+            del self.pipe
+            gc.collect()
+            cuda.empty_cache()
+            destroy_distributed_environment()
+            gc.collect()
 
-        try:
-            # Remove nested attributes if present
-            if hasattr(self.pipe, "llm_engine") and hasattr(self.pipe.llm_engine, "model_executor"):
-                del self.pipe.llm_engine.model_executor
-        except Exception:
-            pass
-
-        del self.pipe
-        del self.tokenizer
-
-        self.pipe = None
-        self.tokenizer = None
-
-        cuda.empty_cache()
-        gc.collect()
+            self.pipe = None
 
     def _process_batch(
         self,
@@ -598,7 +590,14 @@ class Annotator:
             shutil.rmtree(pdout)
         pdout.mkdir(exist_ok=True, parents=True)
 
-        self._load_tokenizer()
+        # We need to clear the model before doing self._load_dataset because the model
+        # cannot be be pickled (which is needed for multiprocessing in dataset.map)
+        if self.pipe is not None and self.num_proc is not None:
+            self.destroy_model()
+
+        if self.tokenizer is None:
+            self._load_tokenizer()
+
         dataset, processed_n_samples = self._load_dataset(
             prompt_template=prompt_template,
             idx_column=idx_column,

llm_annotator-0.2.2/Makefile

@@ -1,11 +0,0 @@
-quality:
-	ruff check src/llm_annotator tests/
-	ruff format --check src/llm_annotator tests/
-
-style:
-	ruff check src/llm_annotator tests/ --fix
-	ruff format src/llm_annotator tests/
-
-setup:
-	uv sync --dev
-	pre-commit install --hook-type pre-push