SDK-43: Add LLM unlearning models support

.github/workflows/lint.yaml

@@ -27,7 +27,7 @@ jobs:
           python -m pip install --upgrade pip
           python -m venv .venv
           source .venv/bin/activate
-          pip install -r requirements/dev.txt -r requirements/pandas.txt -r requirements/polars.txt
+          pip install -r requirements/dev.txt -r requirements/pandas.txt -r requirements/polars.txt -r requirements/transformers.txt
       - run: echo "$PWD/.venv/bin" >> $GITHUB_PATH
       - uses: astral-sh/ruff-action@v3
       - run: ruff check
@@ -50,6 +50,6 @@ jobs:
           python -m pip install --upgrade pip
           python -m venv .venv
           source .venv/bin/activate
-          pip install -r requirements/dev.txt -r requirements/pandas.txt -r requirements/polars.txt
+          pip install -r requirements/dev.txt -r requirements/pandas.txt -r requirements/polars.txt -r requirements/transformers.txt
       - run: echo "$PWD/.venv/bin" >> $GITHUB_PATH
       - run: basedpyright

.github/workflows/pytest-full.yaml

@@ -38,7 +38,7 @@ jobs:
           python -m pip install --upgrade pip
           python -m venv .venv
           source .venv/bin/activate
-          pip install -r requirements/dev.txt -r requirements/polars.txt
+          pip install -r requirements/dev.txt -r requirements/polars.txt -r requirements/transformers.txt
       - name: Run PyTest
         run: .venv/bin/pytest tests/${{ matrix.data-qa-test['test'] }}
         env:

.github/workflows/pytest-sanity.yaml

@@ -56,20 +56,20 @@ jobs:
           python -m pip install --upgrade pip
           python -m venv .venv
           source .venv/bin/activate
-          pip install -r requirements/dev.txt -r requirements/polars.txt
+          pip install -r requirements/dev.txt -r requirements/polars.txt -r requirements/transformers.txt
       - name: Run commands on Windows
         if: github.event_name != 'pull_request' && runner.os == 'Windows' && steps.changes.outputs.non_workflow == 'true'
         run: |
           python -m pip install --upgrade 'pip>=24.1.2'
           python -m venv .venv
           .venv\Scripts\activate
-          python -m pip install -r requirements\dev.txt -r requirements\polars.txt
+          python -m pip install -r requirements\dev.txt -r requirements\polars.txt -r requirements\transformers.txt
       - name: Run PyTest on Linux and macOS
         if: github.event_name != 'pull_request' && runner.os != 'Windows' && steps.changes.outputs.non_workflow == 'true'
         run: .venv/bin/pytest
         env:
-          API_HOST: ${{ secrets.API_HOST }}
-          API_KEY: ${{ secrets.API_KEY }}
+          API_HOST: ${{ secrets.API_HOST2 }}
+          API_KEY: ${{ secrets.API_KEY2 }}
           GCP_CREDENTIALS: ${{ secrets.GCP_CREDENTIALS }}
           AWS_ACCESS_KEY: ${{ secrets.AWS_ACCESS_KEY }}
           AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
@@ -80,8 +80,8 @@ jobs:
         if: github.event_name != 'pull_request' && runner.os == 'Windows' && steps.changes.outputs.non_workflow == 'true'
         run: .venv/Scripts/pytest
         env:
-          API_HOST: ${{ secrets.API_HOST }}
-          API_KEY: ${{ secrets.API_KEY }}
+          API_HOST: ${{ secrets.API_HOST2 }}
+          API_KEY: ${{ secrets.API_KEY2 }}
           GCP_CREDENTIALS: ${{ secrets.GCP_CREDENTIALS }}
           AWS_ACCESS_KEY: ${{ secrets.AWS_ACCESS_KEY }}
           AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}

.github/workflows/vulnerability-scan.yml

@@ -24,6 +24,7 @@ jobs:
           - requirements/docs.txt
           - requirements/pandas.txt
           - requirements/polars.txt
+          - requirements/transformers.txt
     runs-on: ubuntu-latest
     permissions:
       contents: read

.pre-commit-config.yaml

@@ -61,11 +61,18 @@ repos:
       always_run: false
       files: pyproject.toml$
       additional_dependencies: [uv]
-    - id: pip-sync
-      name: sync
+    - id: pip-compile-transformers
+      name: compile requirements/transformers.txt
       entry: uv
-      args: ["pip", "sync", "requirements/dev.txt", "requirements/docs.txt", "requirements/pandas.txt", "requirements/polars.txt"]
+      args: ["pip", "compile", "--extra", "transformers", "-o", "requirements/transformers.txt", "-c", "requirements/requirements.txt"]
       language: python
       always_run: false
-      files: requirements.txt$
+      files: pyproject.toml$
       additional_dependencies: [uv]
+  - repo: https://github.com/astral-sh/uv-pre-commit
+    # uv version.
+    rev: 0.9.6
+    hooks:
+      - id: uv-lock
+      - id: uv-sync
+        args: ["--extra", "dev", "--extra", "docs", "--extra", "pandas", "--extra", "polars", "--extra", "transformers"]

DEV_README.md

@@ -44,12 +44,18 @@ uv pip compile --extra dev -o requirements/dev.txt -c requirements.txt pyproject
 uv pip compile --extra pandas -o requirements/pandas.txt -c requirements.txt pyproject.toml
 uv pip compile --extra polars -o requirements/polars.txt -c requirements.txt pyproject.toml
 uv pip compile --extra docs -o requirements/docs.txt -c requirements.txt pyproject.toml
+uv pip compile --extra transformers -o requirements/transformers.txt -c requirements.txt pyproject.toml
 ```
 
 #### Sync installed packages
 
 ```bash
-uv pip sync requirements/dev.txt requirements/polars.txt
+uv pip sync requirements/dev.txt requirements/pandas.txt requirements/polars.txt requirements/docs.txt requirements/transformers.txt
+```
+or
+
+```bash
+uv sync --extra dev --extra pandas --extra polars --extra docs --extra transformers
 ```
 
 ### Build process

README.md

@@ -67,7 +67,31 @@ You can install the codebase with a simple `pip install hirundo` to install the
 
 ## Usage
 
-Classification example:
+### Unlearning LLM behavior
+
+Make sure to install the `transformers` extra, i.e. `pip install hirundo[transformers]` or `uv pip install hirundo[transformers]` if you have `uv` installed which is much faster than `pip`.
+
+```python
+llm = LlmModel(
+    model_name="Nemotron-Flash-1B",
+    model_source=HuggingFaceTransformersModel(
+        model_name="nvidia/Nemotron-Flash-1B",
+    ),
+)
+llm_id = llm.create()
+run_info = BiasRunInfo(
+    bias_type=BiasType.ALL,
+)
+run_id = LlmUnlearningRun.launch(
+    llm_id,
+    run_info,
+)
+new_adapter = llm.get_hf_pipeline_for_run(run_id)
+```
+
+### Dataset QA
+
+#### Classification example:
 
 ```python
 from hirundo import (
@@ -104,7 +128,7 @@ results = test_dataset.check_run()
 print(results)
 ```
 
-Object detection example:
+#### Object detection example:
 
 ```python
 from hirundo import (

dev.Dockerfile

@@ -7,6 +7,7 @@ COPY . .
 RUN pip install -r requirements/requirements.txt \
     -r requirements/dev.txt -r requirements/docs.txt \
     -r requirements/pandas.txt -r requirements/polars.txt \
+    -r requirements/transformers.txt \
      && pip install ipykernel
 
 CMD ["python"]

hirundo/__init__.py

@@ -30,6 +30,15 @@
     StorageGit,
     StorageS3,
 )
+from .unlearning_llm import (
+    BiasRunInfo,
+    BiasType,
+    HuggingFaceTransformersModel,
+    LlmModel,
+    LlmSources,
+    LlmUnlearningRun,
+    LocalTransformersModel,
+)
 from .unzip import load_df, load_from_zip
 
 __all__ = [
@@ -59,8 +68,15 @@
     "StorageGit",
     "StorageConfig",
     "DatasetQAResults",
+    "BiasRunInfo",
+    "BiasType",
+    "HuggingFaceTransformersModel",
+    "LlmModel",
+    "LlmSources",
+    "LlmUnlearningRun",
+    "LocalTransformersModel",
     "load_df",
     "load_from_zip",
 ]
 
-__version__ = "0.1.21"
+__version__ = "0.1.22"

hirundo/_llm_pipeline.py

@@ -0,0 +1,153 @@
+import importlib.util
+import tempfile
+import zipfile
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+from hirundo import HirundoError
+from hirundo._http import requests
+from hirundo._timeouts import DOWNLOAD_READ_TIMEOUT
+from hirundo.logger import get_logger
+
+if TYPE_CHECKING:
+    from torch import device as torch_device
+    from transformers.configuration_utils import PretrainedConfig
+    from transformers.modeling_utils import PreTrainedModel
+    from transformers.pipelines.base import Pipeline
+
+    from hirundo.unlearning_llm import LlmModel, LlmModelOut
+
+logger = get_logger(__name__)
+
+
+ZIP_FILE_CHUNK_SIZE = 50 * 1024 * 1024  # 50 MB
+REQUIRED_PACKAGES_FOR_PIPELINE = ["peft", "transformers", "accelerate"]
+
+
+def get_hf_pipeline_for_run_given_model(
+    llm: "LlmModel | LlmModelOut",
+    run_id: str,
+    config: "PretrainedConfig | None" = None,
+    device: "str | int | torch_device | None" = None,
+    device_map: str | dict[str, int | str] | None = None,
+    trust_remote_code: bool = False,
+    token: str | None = None,
+) -> "Pipeline":
+    for package in REQUIRED_PACKAGES_FOR_PIPELINE:
+        if importlib.util.find_spec(package) is None:
+            raise HirundoError(
+                f'{package} is not installed. Please install transformers extra with pip install "hirundo[transformers]"'
+            )
+    from peft import PeftModel
+    from transformers.models.auto.configuration_auto import AutoConfig
+    from transformers.models.auto.modeling_auto import (
+        MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES,
+        AutoModelForCausalLM,
+        AutoModelForImageTextToText,
+    )
+    from transformers.models.auto.tokenization_auto import AutoTokenizer
+    from transformers.pipelines import pipeline
+
+    from hirundo.unlearning_llm import (
+        HuggingFaceTransformersModel,
+        HuggingFaceTransformersModelOutput,
+        LlmUnlearningRun,
+    )
+
+    run_results = LlmUnlearningRun.check_run_by_id(run_id)
+    if run_results is None:
+        raise HirundoError("No run results found")
+    result_payload = (
+        run_results.get("result", run_results)
+        if isinstance(run_results, dict)
+        else run_results
+    )
+    if isinstance(result_payload, dict):
+        result_url = result_payload.get("result")
+    else:
+        result_url = result_payload
+    if not isinstance(result_url, str):
+        raise HirundoError("Run results did not include a download URL")
+    # Stream the zip file download
+
+    zip_file_path = tempfile.NamedTemporaryFile(delete=False).name
+    with requests.get(
+        result_url,
+        timeout=DOWNLOAD_READ_TIMEOUT,
+        stream=True,
+    ) as r:
+        r.raise_for_status()
+        with open(zip_file_path, "wb") as zip_file:
+            for chunk in r.iter_content(chunk_size=ZIP_FILE_CHUNK_SIZE):
+                zip_file.write(chunk)
+        logger.info(
+            "Successfully downloaded the result zip file for run ID %s to %s",
+            run_id,
+            zip_file_path,
+        )
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_dir_path = Path(temp_dir)
+        with zipfile.ZipFile(zip_file_path, "r") as zip_file:
+            zip_file.extractall(temp_dir_path)
+        # Attempt to load the tokenizer normally
+        base_model_name = (
+            llm.model_source.model_name
+            if isinstance(
+                llm.model_source,
+                HuggingFaceTransformersModel | HuggingFaceTransformersModelOutput,
+            )
+            else llm.model_source.local_path
+        )
+        token = (
+            llm.model_source.token
+            if isinstance(
+                llm.model_source,
+                HuggingFaceTransformersModel,
+            )
+            else token
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            base_model_name,
+            token=token,
+            trust_remote_code=trust_remote_code,
+        )
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        config = AutoConfig.from_pretrained(
+            base_model_name,
+            token=token,
+            trust_remote_code=trust_remote_code,
+        )
+        config_dict = config.to_dict() if hasattr(config, "to_dict") else config
+        is_multimodal = (
+            config_dict.get("model_type")
+            in MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.keys()
+        )
+        if is_multimodal:
+            base_model = AutoModelForImageTextToText.from_pretrained(
+                base_model_name,
+                token=token,
+                trust_remote_code=trust_remote_code,
+            )
+        else:
+            base_model = AutoModelForCausalLM.from_pretrained(
+                base_model_name,
+                token=token,
+                trust_remote_code=trust_remote_code,
+            )
+        model = cast(
+            "PreTrainedModel",
+            PeftModel.from_pretrained(
+                base_model, str(temp_dir_path / "unlearned_model_folder")
+            ),
+        )
+
+        return pipeline(
+            task="text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            config=config,
+            device=device,
+            device_map=device_map,
+        )