transformers-knn-adapter 0.2.0

Hugging Face Transformers image embedding adapter with a scikit-learn KNN classification head.

transformers-knn-adapter

transformers_knn_adapter extends Hugging Face image models by attaching a scikit-learn KNN classifier on top of transformer embeddings.

Requirements

• Python 3.11+
• uv

Setup

uv sync --dev

Run Tests

uv run pytest

Train

uv run python -m transformers_knn_adapter.knn_image_pipeline train \
  --model microsoft/resnet-50 \
  --knn-model-path /tmp/knn/dinov2_small_mini_imagenet_full.joblib \
  --dataset timm/mini-imagenet \
  --split train \
  --max-samples 1000 \
  --shuffle \
  --grid-search \
  --grid-search-splits 3 \
  --grid-search-repeats 2 \
  --grid-search-scoring f1_macro

Evaluate

uv run python -m transformers_knn_adapter.knn_image_pipeline eval \
  --model microsoft/resnet-50 \
  --knn-model-path /tmp/knn/dinov2_small_mini_imagenet_full.joblib \
  --dataset timm/mini-imagenet \
  --split test \
  --stratified \
  --max-samples 100 \
  --shuffle \
  --batch-size 100

Inference

uv run python -m transformers_knn_adapter.knn_image_pipeline infer \
  --model microsoft/resnet-50 \
  --knn-model-path /tmp/knn/dinov2_small_mini_imagenet_full.joblib \
  --image https://picsum.photos/200 \
  --inference-batch-size 5

Python API

from transformers_knn_adapter.knn_image_pipeline import pipeline

clf = pipeline(
    "image-classification",
    model_path="microsoft/resnet-50",
    knn_model_path="/tmp/knn/model.joblib",
)

Notes

Real train/eval runs can download model and dataset artifacts from Hugging Face.