Update formatting

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 23.9.1
+    rev: 24.1.1
     hooks:
       - id: black
         language_version: python3.11

setup.py

@@ -22,13 +22,13 @@
 requirements_huggingface = [
     "sentencepiece>=0.1.99",  # for Llama tokenizers. cappr should work out-of-the-box
     "torch>=1.12.1",
-    "transformers>=4.31.0",  # high version b/c Llama
+    "transformers>=4.31.0",
 ]
 
 requirements_huggingface_dev = [
     req if not req.startswith("transformers>=") else "transformers>=4.35.0"
-    # To test Mistral in our testing workflow, we need >=4.34.0. To demo AutoGPTQ on CPU
-    # and AutoAWQ with caching, need >=4.35.0
+    # To test Mistral in our testing workflow, we need >=4.34.0.
+    # To demo AutoGPTQ on CPU and AutoAWQ with caching, need >=4.35.0.
     for req in requirements_huggingface
 ] + ["huggingface-hub>=0.16.4"]
 

src/cappr/__init__.py

@@ -3,6 +3,7 @@
 
 https://cappr.readthedocs.io/
 """
+
 __version__ = "0.8.8"
 
 from . import utils

src/cappr/huggingface/__init__.py

@@ -7,4 +7,5 @@
 
 https://cappr.readthedocs.io/en/latest/select_a_language_model.html#huggingface
 """
+
 from . import _utils, classify, classify_no_cache

src/cappr/huggingface/_utils.py

@@ -1,6 +1,7 @@
 """
 YouTils
 """
+
 from __future__ import annotations
 from contextlib import contextmanager, ExitStack, nullcontext
 from functools import lru_cache
@@ -59,14 +60,21 @@ def _no_grad(model: ModelForCausalLM):  # model given to keep interface the same
         yield
 
 
+# Some models don't perfectly implement the HF model call interface. In particular,
+# they're missing the return_dict and use_cache kwargs. They're instead in the model
+# config. I see that as a more extensible design anyway.
+
+
 @contextmanager
 def _return_dict(model: ModelForCausalLM):
     """
     In this context, the model returns a dataclass when it's called.
     """
-    with _setattr(model.config, "return_dict", True) if hasattr(
-        model, "config"
-    ) else nullcontext():  # null b/c just try model(...).logits when needed
+    with (
+        _setattr(model.config, "return_dict", True)
+        if hasattr(model, "config")
+        else nullcontext()
+    ):  # null b/c just try model(...).logits when needed
         yield
 
 
@@ -75,9 +83,11 @@ def _use_cache(model: ModelForCausalLM):
     """
     In this context, the model output includes a `past_key_values` attribute.
     """
-    with _setattr(model.config, "use_cache", True) if hasattr(
-        model, "config"
-    ) else nullcontext():  # null b/c just try model(...).past_key_values when needed
+    with (
+        _setattr(model.config, "use_cache", True)
+        if hasattr(model, "config")
+        else nullcontext()
+    ):  # null b/c just try model(...).past_key_values when needed
         yield
 
 

src/cappr/huggingface/classify.py

@@ -9,6 +9,7 @@
 In the implementation, attention block keys and values for prompts are automatically
 cached and shared across completions.
 """
+
 from __future__ import annotations
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass

src/cappr/huggingface/classify_no_cache.py

@@ -10,6 +10,7 @@
 This module is a mirror of :mod:`cappr.huggingface.classify`. The difference is that
 this module **does not** cache attention keys and values.
 """
+
 from __future__ import annotations
 from typing import cast, Literal, Mapping, Sequence
 

src/cappr/llama_cpp/__init__.py

@@ -7,4 +7,5 @@
 
 https://cappr.readthedocs.io/en/latest/select_a_language_model.html#llama-cpp
 """
+
 from . import _utils, classify, _classify_no_cache

src/cappr/llama_cpp/_classify_no_cache.py

@@ -1,6 +1,7 @@
 """
 Mirror. For testing purposes only. Is strictly slower
 """
+
 from __future__ import annotations
 from functools import partial
 from typing import Literal, Sequence

src/cappr/llama_cpp/_utils.py

@@ -1,6 +1,7 @@
 """
 Utilz
 """
+
 from __future__ import annotations
 from contextlib import contextmanager
 from functools import lru_cache

src/cappr/llama_cpp/classify.py

@@ -19,6 +19,7 @@
     --local-dir . \\
     --local-dir-use-symlinks False
 """
+
 from __future__ import annotations
 from contextlib import contextmanager
 from typing import cast, Literal, Sequence

src/cappr/openai/__init__.py

@@ -3,4 +3,5 @@
 
 https://cappr.readthedocs.io/en/latest/select_a_language_model.html#openai
 """
+
 from . import api, classify

src/cappr/openai/classify.py

@@ -5,6 +5,7 @@
 
 You probably just want the :func:`predict` or :func:`predict_examples` functions :-)
 """
+
 from __future__ import annotations
 from typing import Literal, Sequence
 

src/cappr/utils/__init__.py

@@ -1,4 +1,5 @@
 """
 Utilities shared across model backends
 """
+
 from . import _batch, _check, _no_cache, classify

src/cappr/utils/_batch.py

@@ -1,6 +1,7 @@
 """
 Batch lists into sublists of constant or variable sizes, and batchify functions
 """
+
 from __future__ import annotations
 from functools import wraps
 import inspect

src/cappr/utils/_check.py

@@ -10,6 +10,7 @@
     - inputting an incorrectly structured prior will cause all of the model's compute to
       be a waste.
 """
+
 from __future__ import annotations
 from typing import Callable, Literal, Sequence
 

src/cappr/utils/_no_cache.py

@@ -1,6 +1,7 @@
 """
 Generic implementation for non-cached LM interfaces
 """
+
 from __future__ import annotations
 from functools import lru_cache
 from typing import Any, Callable, cast, Literal, Sequence

src/cappr/utils/classify.py

@@ -2,6 +2,7 @@
 Transform completion token log-probabilites into a probability distribution over
 completions
 """
+
 from __future__ import annotations
 from contextlib import contextmanager
 from functools import wraps
@@ -555,9 +556,9 @@ def _predict_examples(predict_proba_examples_func):
     def wrapper(
         examples: Example | Sequence[Example], *args, **kwargs
     ) -> str | list[str]:
-        pred_probs: npt.NDArray[np.floating] | list[
-            npt.NDArray[np.floating]
-        ] = predict_proba_examples_func(examples, *args, **kwargs)
+        pred_probs: npt.NDArray[np.floating] | list[npt.NDArray[np.floating]] = (
+            predict_proba_examples_func(examples, *args, **kwargs)
+        )
         if isinstance(examples, Example):
             # User convenience: examples is a singleton
             assert pred_probs.ndim == 1  # double check

tests/_base.py

@@ -2,6 +2,7 @@
 Base classes which are parametrized with a set of test cases which every `classify`
 module must pass.
 """
+
 from __future__ import annotations
 from typing import Sequence
 

tests/_test_content.py

@@ -2,6 +2,7 @@
 Helper functions to test that outputs of `classify` modules have the correct content,
 assuming a reference/correct implementation exists.
 """
+
 from __future__ import annotations
 from typing import Sequence
 

tests/_test_form.py

@@ -3,6 +3,7 @@
   - outputs have the correct shape and form (content is NOT checked)
   - inputs are correctly checked.
 """
+
 from __future__ import annotations
 import re
 from typing import Sequence

tests/huggingface/test_huggingface_classify.py

@@ -3,6 +3,7 @@
 functions' outputs are numerically close to those from
 `cappr.huggingface.classify_no_cache`.
 """
+
 from __future__ import annotations
 from contextlib import nullcontext
 import os

tests/llama_cpp/test_llama_cpp_classify.py

@@ -3,6 +3,7 @@
 functions' outputs are numerically close to those from
 `cappr.llama_cpp._classify_no_cache`.
 """
+
 from __future__ import annotations
 from dataclasses import dataclass
 import os

tests/openai/test_openai_api.py

@@ -1,6 +1,7 @@
 """
 Unit tests `cappr.openai.api`. Currently pretty barebones.
 """
+
 from __future__ import annotations
 import os
 

tests/openai/test_openai_classify.py

@@ -1,6 +1,7 @@
 """
 Unit and integration tests for tests `cappr.openai.classify`.
 """
+
 from __future__ import annotations
 import os
 import sys

tests/test_Example.py

@@ -1,6 +1,7 @@
 """
 Unit tests `cappr.Example` input checks.
 """
+
 from __future__ import annotations
 import re
 from typing import Any, Sequence

tests/utils/test_utils__batch.py

@@ -1,6 +1,7 @@
 """
 Unit tests `cappr.utils._batch`.
 """
+
 from __future__ import annotations
 import re
 

tests/utils/test_utils_classify.py

@@ -1,6 +1,7 @@
 """
 Unit tests `cappr.utils.classify`.
 """
+
 from __future__ import annotations
 from typing import Any