migate to uv & ruff, format notebooks (#56)

* migrate to uv and ruff from poetry and black * remove accelerate from dependencies. please install accelerate manually if you need `device_map="auto"` * raise newest supported python to `3.13` * test python `3.12` and `3.13` in CI this pr also rolled in a small tweak to the `repeng.saes.from_eleuther` interface
vgel · Dec 14, 2024 · e205758 · e205758
1 parent 021022a
commit e205758
Show file tree

Hide file tree

Showing 14 changed files with 1,458 additions and 1,993 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -10,31 +10,50 @@ jobs:
   format:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: run black
-      # make sure to keep this in sync with pyproject.toml
-      run: pipx run -- "black>=24.2.0,<25.0.0" --check .
+    - uses: actions/checkout@v4
+
+    - name: install uv
+      uses: astral-sh/setup-uv@v4
+      with:
+        version: "0.5.9"
+
+    - name: ruff check
+      run: uv run ruff check
+
+    - name: ruff format --check
+      run: uv run ruff format --check
 
   test:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10", "3.11"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
 
     steps:
-    - uses: actions/checkout@v2
-    - run: pipx install poetry
-    - name: install python ${{ matrix.python-version }}
+    - uses: actions/checkout@v4
+
+    # don't use uv install python because we need a system python for
+    # numpy wheel builds
+    - name: install python
       uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
-        cache: 'poetry'
-    - run: poetry install
+
+    - name: install uv
+      uses: astral-sh/setup-uv@v4
+      with:
+        version: "0.5.9"
+        enable-cache: true
+        cache-dependency-glob: "uv.lock"
+
+    - run: uv sync --all-extras --dev
+
     - name: cache huggingface
       uses: actions/cache@v4
       with:
         path: ~/.cache/huggingface
         key: ${{ runner.os }}-hf-models-${{ hashFiles('**/lockfiles') }}
         restore-keys: |
           ${{ runner.os }}-hf-models-
-    - run: poetry run pytest
+
+    - run: uv run pytest
diff --git a/.gitignore b/.gitignore
@@ -1,7 +1,8 @@
 # project
+__pycache__/
+.venv/
 
 # system
-__pycache__/
 *.pyc
 *.pyo
 *.pyd

diff --git a/notebooks/emotion.ipynb b/notebooks/emotion.ipynb
@@ -54,7 +54,13 @@
     "tokenizer.pad_token_id = 0\n",
     "\n",
     "model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)\n",
-    "model = model.to(\"cuda:0\" if torch.cuda.is_available() else \"mps:0\" if torch.backends.mps.is_available() else \"cpu\")\n",
+    "model = model.to(\n",
+    "    \"cuda:0\"\n",
+    "    if torch.cuda.is_available()\n",
+    "    else \"mps:0\"\n",
+    "    if torch.backends.mps.is_available()\n",
+    "    else \"cpu\"\n",
+    ")\n",
     "model = ControlModel(model, list(range(-5, -18, -1)))\n",
     "\n",
     "user_tag, asst_tag = \"[INST]\", \"[/INST]\""
@@ -89,15 +95,20 @@
     "# good: \"high on acid\" / \"sober, not on acid\" — the negative prompt is more directly opposite\n",
     "positive_personas = [\"happy\", \"ecstatic\", \"delighted\"]\n",
     "negative_personas = [\"sad\", \"depressed\", \"dismayed\"]\n",
+    "\n",
+    "\n",
     "def template(persona: str, suffix: str) -> str:\n",
     "    return f\"{user_tag} Act as if you're extremely {persona}. {asst_tag} {suffix}\"\n",
     "\n",
+    "\n",
     "dataset = []\n",
     "for suffix in suffixes:\n",
     "    tokens = tokenizer.tokenize(suffix)\n",
     "    for i in range(1, len(tokens)):\n",
     "        truncated = tokenizer.convert_tokens_to_string(tokens[:i])\n",
-    "        for positive_persona, negative_persona in zip(positive_personas, negative_personas):\n",
+    "        for positive_persona, negative_persona in zip(\n",
+    "            positive_personas, negative_personas\n",
+    "        ):\n",
     "            dataset.append(\n",
     "                DatasetEntry(\n",
     "                    positive=template(positive_persona, truncated),\n",
@@ -127,7 +138,7 @@
     }
    ],
    "source": [
-    "model.reset() # make sure you always reset the model before training a new vector\n",
+    "model.reset()  # make sure you always reset the model before training a new vector\n",
     "control_vector = ControlVector.train(\n",
     "    model,\n",
     "    tokenizer,\n",
@@ -164,10 +175,10 @@
     "# tokenizer and generation settings\n",
     "input_ids = tokenizer(input, return_tensors=\"pt\").to(model.device)\n",
     "settings = {\n",
-    "    \"pad_token_id\": tokenizer.eos_token_id, # silence warning\n",
-    "    \"do_sample\": False, # temperature=0\n",
+    "    \"pad_token_id\": tokenizer.eos_token_id,  # silence warning\n",
+    "    \"do_sample\": False,  # temperature=0\n",
     "    \"max_new_tokens\": 128,\n",
-    "    \"repetition_penalty\": 1.1, # reduce control jank\n",
+    "    \"repetition_penalty\": 1.1,  # reduce control jank\n",
     "}\n",
     "\n",
     "print(\"==baseline\")\n",

diff --git a/notebooks/experiments.ipynb b/notebooks/experiments.ipynb
@@ -54,7 +54,13 @@
     "tokenizer.pad_token_id = 0\n",
     "\n",
     "model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)\n",
-    "model = model.to(\"cuda:0\" if torch.cuda.is_available() else \"mps:0\" if torch.backends.mps.is_available() else \"cpu\")\n",
+    "model = model.to(\n",
+    "    \"cuda:0\"\n",
+    "    if torch.cuda.is_available()\n",
+    "    else \"mps:0\"\n",
+    "    if torch.backends.mps.is_available()\n",
+    "    else \"cpu\"\n",
+    ")\n",
     "model = ControlModel(model, list(range(-5, -18, -1)))\n",
     "\n",
     "user_tag, asst_tag = \"[INST]\", \"[/INST]\""
@@ -88,15 +94,18 @@
     "    for i in range(1, len(tokens) - 5)\n",
     "]\n",
     "\n",
+    "\n",
     "def make_dataset(\n",
     "    template: str,\n",
     "    positive_personas: list[str],\n",
     "    negative_personas: list[str],\n",
-    "    suffix_list: list[str]\n",
+    "    suffix_list: list[str],\n",
     ") -> list[DatasetEntry]:\n",
     "    dataset = []\n",
     "    for suffix in suffix_list:\n",
-    "        for positive_persona, negative_persona in zip(positive_personas, negative_personas):\n",
+    "        for positive_persona, negative_persona in zip(\n",
+    "            positive_personas, negative_personas\n",
+    "        ):\n",
     "            positive_template = template.format(persona=positive_persona)\n",
     "            negative_template = template.format(persona=negative_persona)\n",
     "            dataset.append(\n",
@@ -131,21 +140,23 @@
     "        input = f\"{user_tag} {input.strip()} {asst_tag}\"\n",
     "    input_ids = tokenizer(input, return_tensors=\"pt\").to(model.device)\n",
     "    settings = {\n",
-    "        \"pad_token_id\": tokenizer.eos_token_id, # silence warning\n",
-    "        \"do_sample\": False, # temperature=0\n",
+    "        \"pad_token_id\": tokenizer.eos_token_id,  # silence warning\n",
+    "        \"do_sample\": False,  # temperature=0\n",
     "        \"max_new_tokens\": max_new_tokens,\n",
     "        \"repetition_penalty\": repetition_penalty,\n",
     "    }\n",
     "\n",
     "    if show_baseline:\n",
     "        print(\"==baseline ---------------------------------------------------\")\n",
     "        model.reset()\n",
-    "        print(tokenizer.decode(model.generate(**input_ids, **settings).squeeze()).strip())\n",
-    "    \n",
+    "        print(\n",
+    "            tokenizer.decode(model.generate(**input_ids, **settings).squeeze()).strip()\n",
+    "        )\n",
+    "\n",
     "    print(\"\\n++control ---------------------------------------------------\")\n",
     "    model.set_control(vector, positive_coeff)\n",
     "    print(tokenizer.decode(model.generate(**input_ids, **settings).squeeze()).strip())\n",
-    "    \n",
+    "\n",
     "    print(\"\\n--control ---------------------------------------------------\")\n",
     "    model.set_control(vector, negative_coeff)\n",
     "    print(tokenizer.decode(model.generate(**input_ids, **settings).squeeze()).strip())\n",
@@ -286,7 +297,7 @@
     "    \"Act as if you're extremely {persona}.\",\n",
     "    [\"high on psychedelic drugs\"],\n",
     "    [\"sober from psychedelic drugs\"],\n",
-    "    truncated_output_suffixes_512, # gives (subjectively) better results with slightly fewer samples\n",
+    "    truncated_output_suffixes_512,  # gives (subjectively) better results with slightly fewer samples\n",
     ")\n",
     "model.reset()\n",
     "trippy_vector = ControlVector.train(model, tokenizer, trippy_dataset)"
@@ -656,7 +667,7 @@
    ],
    "source": [
     "generate_with_vector(\n",
-    "    f\"Tell me a recent scientific breakthrough.\",\n",
+    "    \"Tell me a recent scientific breakthrough.\",\n",
     "    future_vector,\n",
     "    (2.1, -2),\n",
     ")"
@@ -793,7 +804,9 @@
     "dealership_dataset = make_dataset(\n",
     "    \"Pretend you're {persona} making statements about the world.\",\n",
     "    [\"an assistant for a car dealersip, only interested in cars and the dealership,\"],\n",
-    "    [\"a random person, who talks about anything and doesn't care about cars or the dealership,\"],\n",
+    "    [\n",
+    "        \"a random person, who talks about anything and doesn't care about cars or the dealership,\"\n",
+    "    ],\n",
     "    truncated_fact_suffixes,\n",
     ")\n",
     "model.reset()\n",

diff --git a/notebooks/honesty.ipynb b/notebooks/honesty.ipynb
@@ -54,7 +54,13 @@
     "tokenizer.pad_token_id = 0\n",
     "\n",
     "model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)\n",
-    "model = model.to(\"cuda:0\" if torch.cuda.is_available() else \"mps:0\" if torch.backends.mps.is_available() else \"cpu\")\n",
+    "model = model.to(\n",
+    "    \"cuda:0\"\n",
+    "    if torch.cuda.is_available()\n",
+    "    else \"mps:0\"\n",
+    "    if torch.backends.mps.is_available()\n",
+    "    else \"cpu\"\n",
+    ")\n",
     "model = ControlModel(model, list(range(-5, -18, -1)))\n",
     "\n",
     "user_tag, asst_tag = \"[INST]\", \"[/INST]\""
@@ -86,15 +92,20 @@
     "# see the emotion notebook for more details about this\n",
     "positive_personas = [\"honest\"]\n",
     "negative_personas = [\"untruthful\"]\n",
+    "\n",
+    "\n",
     "def template(persona: str, suffix: str) -> str:\n",
     "    return f\"{user_tag} Pretend you're an {persona} person making statements about the world. {asst_tag} {suffix}\"\n",
     "\n",
+    "\n",
     "dataset = []\n",
     "for suffix in suffixes:\n",
     "    tokens = tokenizer.tokenize(suffix)\n",
     "    for i in range(1, len(tokens) - 5):\n",
     "        truncated = tokenizer.convert_tokens_to_string(tokens[:i])\n",
-    "        for positive_persona, negative_persona in zip(positive_personas, negative_personas):\n",
+    "        for positive_persona, negative_persona in zip(\n",
+    "            positive_personas, negative_personas\n",
+    "        ):\n",
     "            dataset.append(\n",
     "                DatasetEntry(\n",
     "                    positive=template(positive_persona, truncated),\n",
@@ -124,7 +135,7 @@
     }
    ],
    "source": [
-    "model.reset() # make sure you always reset the model before training a new vector\n",
+    "model.reset()  # make sure you always reset the model before training a new vector\n",
     "control_vector = ControlVector.train(\n",
     "    model,\n",
     "    tokenizer,\n",
@@ -163,10 +174,10 @@
     "# tokenizer and generation settings\n",
     "input_ids = tokenizer(input, return_tensors=\"pt\").to(model.device)\n",
     "settings = {\n",
-    "    \"pad_token_id\": tokenizer.eos_token_id, # silence warning\n",
-    "    \"do_sample\": False, # temperature=0\n",
+    "    \"pad_token_id\": tokenizer.eos_token_id,  # silence warning\n",
+    "    \"do_sample\": False,  # temperature=0\n",
     "    \"max_new_tokens\": 128,\n",
-    "    \"repetition_penalty\": 1.1, # reduce control jank\n",
+    "    \"repetition_penalty\": 1.1,  # reduce control jank\n",
     "}\n",
     "\n",
     "print(\"==baseline\")\n",