Migration to uv (#108)

Signed-off-by: Max Jeblick <maximilianjeblick@gmail.com>

Author: alessiodevoto

.flake8

@@ -1,5 +1,5 @@
 [flake8]
-exclude = .venv,venv,.git,__pycache__,build,dist, .mypy_cache
+exclude = .venv,venv,.git,__pycache__,build,dist,.mypy_cache,.pytest_cache
 max-line-length = 120
 per-file-ignores =
     __init__.py:F401

.github/workflows/python-publish.yml

@@ -18,12 +18,13 @@ jobs:
       uses: actions/setup-python@v3
       with:
         python-version: 3.10.11
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install build
-    - name: Build package
-      run: python -m build
+    
+    - name: Install uv
+      uses: astral-sh/setup-uv@v6
+
+    - name: Build
+      run: uv build --no-sources
+
     - name: Publish package
       uses: pypa/gh-action-pypi-publish@release/v1
       with:

.github/workflows/style.yml

@@ -15,13 +15,13 @@ jobs:
         with:
           python-version: 3.10.11
 
-      - name: Install Poetry
-        run: |
-          curl -sSL https://install.python-poetry.org | python3 -
-          echo "$HOME/.local/bin" >> $GITHUB_PATH  # Add Poetry to the PATH
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
 
       - name: Install dependencies
-        run: poetry install --with dev
+        run: uv sync --all-groups
 
       - name: Run style checks
         run: make style

.github/workflows/test.yml

@@ -2,7 +2,7 @@ name: Test
 
 on:
   push:
-    branches: [ main ]
+      branches: [ main ]
   pull_request:
 
 jobs:
@@ -15,12 +15,12 @@ jobs:
         with:
           python-version: 3.10.11
 
-      - name: Install Poetry
-        run: |
-          curl -sSL https://install.python-poetry.org | python3 -
-          echo "$HOME/.local/bin" >> $GITHUB_PATH  # Add Poetry to the PATH
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          enable-cache: true
 
       - name: Install dependencies
-        run: poetry install --with dev
+        run: uv sync --all-groups
 
       - run: make test

.gitignore

@@ -2,7 +2,7 @@ dev_notebooks/
 results/
 reports/
 .DS_Store
-poetry.lock
+uv.lock
 *.parquet
 # Byte-compiled / optimized / DLL files
 __pycache__/

Makefile

@@ -1,17 +1,17 @@
 SHELL := /bin/bash
-POETRY ?= $(shell which poetry)
+UV ?= $(shell which uv)
 BUILD_VERSION:=$(APP_VERSION)
 TESTS_FILTER:=
 
 PYTEST_LOG=--log-cli-level=debug --log-format="%(asctime)s %(levelname)s [%(name)s:%(filename)s:%(lineno)d] %(message)s" --log-date-format="%Y-%m-%d %H:%M:%S"
 
 .PHONY: isort
 isort:
-	$(POETRY) run isort .
+	$(UV) run isort .
 
 .PHONY: black
 black:
-	$(POETRY) run black .
+	$(UV) run black .
 
 PHONY: format
 format: isort black
@@ -24,10 +24,10 @@ style: reports
 	@echo -n > reports/copyright_errors.log
 	@echo
 
-	-$(POETRY) run flake8 | tee -a reports/flake8_errors.log
+	-$(UV) run flake8 | tee -a reports/flake8_errors.log
 	@if [ -s reports/flake8_errors.log ]; then exit 1; fi
 
-	-$(POETRY) run mypy . --check-untyped-defs | tee -a reports/mypy.log
+	-$(UV) run mypy . --check-untyped-defs | tee -a reports/mypy.log
 	@if ! grep -Eq "Success: no issues found in [0-9]+ source files" reports/mypy.log ; then exit 1; fi
 
 	@echo "Checking for SPDX-FileCopyrightText headers in Python files..."
@@ -42,7 +42,7 @@ reports:
 .PHONY: test
 test: reports
 	PYTHONPATH=. \
-	$(POETRY) run pytest \
+	$(UV) run pytest \
 		--cov-report xml:reports/coverage.xml \
 		--cov=kvpress/ \
 		--junitxml=./reports/junit.xml \

README.md

@@ -16,22 +16,17 @@ Deploying long-context LLMs is costly due to the linear growth of the key-value
 pip install kvpress
 ```
 
-If possible, install flash attention:
-```bash
-pip install flash-attn --no-build-isolation
-```
-
-For a local installation with all dev dependencies, use poetry:
+For a local installation with all dev dependencies, use uv:
 
 ```bash
 git clone https://github.com/NVIDIA/kvpress.git
 cd kvpress
-poetry install --with dev
+uv sync --all-groups
 ```
 
 ## Usage
 
-kvpress provides a set of "presses" that compress the KV cache during the prefilling-phase. Each press is associated with a `compression_ratio` attribute that measures the compression of the cache. The easiest way to use a press is through our custom `KVPressTextGenerationPipeline`. It is automatically registered as a transformers pipeline with the name "kv-press-text-generation" when kvpress is imported and handles chat templates and tokenization for you:
+KVPress provides a set of "presses" that compress the KV cache during the prefilling-phase. Each press is associated with a `compression_ratio` attribute that measures the compression of the cache. The easiest way to use a press is through our custom `KVPressTextGenerationPipeline`. It is automatically registered as a transformers pipeline with the name "kv-press-text-generation" when kvpress is imported and handles chat templates and tokenization for you:
 
 ```python
 from transformers import pipeline
@@ -208,4 +203,25 @@ with press(model):
 
 However, the `generate` method does not allow to exclude the question from the compression, which would artificially favors methods such as SnapKV. Ideally, we want a compression method that works whatever comes after the context (_e.g._ for use cases such as chat or document question answering). Finally the `generate` method does not allow to provide generation for multiple questions at once.
 
-</details>
+</details>
+
+
+## Advanced installation settings
+To install optional packages, you can use [uv](https://docs.astral.sh/uv/). 
+To install with flash attention, just run:
+
+```bash
+git clone https://github.com/NVIDIA/kvpress.git
+cd kvpress
+uv sync --extra flash-attn
+```
+
+To install with dependencies for evaluation, run 
+
+```bash
+git clone https://github.com/NVIDIA/kvpress.git
+cd kvpress
+uv sync --extra eval
+```
+
+Notice that optional dependecies can be combined.

evaluation/README.md

@@ -5,6 +5,7 @@
 We support evaluation for all the presses implemented in the library, on a variety of popular benchmarks.
 
 ### Quick Start 🚀
+> Evaluation requires some additional packages. You can install them with `uv sync --group eval`
 
 Running evaluation is straightforward! Make sure you are in the `evaluation` directory, then:
 

kvpress/pipeline.py

@@ -135,7 +135,10 @@ def preprocess(
         else:
             separator = "\n" + "#" * len(context)
             context = self.tokenizer.apply_chat_template(
-                [{"role": "user", "content": context + separator}], add_generation_prompt=True, tokenize=False
+                [{"role": "user", "content": context + separator}],
+                add_generation_prompt=True,
+                tokenize=False,
+                enable_thinking=False,
             )
             context, question_suffix = context.split(separator)
 

kvpress/presses/block_press.py

@@ -16,10 +16,11 @@ class BlockPress(BasePress):
     BlockPress: Block-wise iterative KV cache compression.
 
     Applies compression in fixed-size blocks. Iteratively scores and prunes tokens block by block, maintaining
-    a buffer of previously kept tokens for context. Mathematically equivalent
-    to global compression when scoring uses only local information.
+    a buffer of previously kept tokens for context. Mathematically equivalent to global compression when
+    scoring uses only local information. It was introduced in the KeyDiff paper as part of the KeyDiff press,
+    but it can also work as a standalone press.
 
-    Based on BlockPress (https://arxiv.org/abs/2504.15364).
+    Based on the KeyDiff paper (https://arxiv.org/abs/2504.15364).
 
     Parameters
     ----------