update skeletoken

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
       - id: end-of-file-fixer
         description: Makes sure files end in a newline and only a newline.
       - id: check-added-large-files
-        args: ['--maxkb=5000']
+        args: ["--maxkb=5000"]
         description: Prevent giant files from being committed.
       - id: check-case-conflict
         description: Check for files with names that would conflict on case-insensitive filesystems like MacOS/Windows.
@@ -20,10 +20,10 @@ repos:
     hooks:
       - id: pydoclint
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.13
+    rev: v0.15.1
     hooks:
       - id: ruff
-        args: [ --fix ]
+        args: [--fix]
       - id: ruff-format
   - repo: local
     hooks:

pynife/data.py

@@ -21,8 +21,7 @@ def _pair_stream(
     txt_path: Path,
     emb_path: Path,
 ) -> Iterator[tuple[dict[str, str], np.ndarray]]:
-    """
-    Stream aligned (text, embedding_row) pairs from a txt jsonl and tensor file.
+    """Stream aligned (text, embedding_row) pairs from a txt jsonl and tensor file.
 
     The text path can contain arbitrary dictionary fields, but will also contain a "text" field.
 
@@ -37,7 +36,7 @@ def _pair_stream(
     embs = torch.load(emb_path)
     embs = embs.float().numpy()
 
-    with open(txt_path, "r", encoding="utf-8") as f:
+    with open(txt_path, encoding="utf-8") as f:
         for i, line in enumerate(f):
             item = json.loads(line)
             yield item, embs[i]
@@ -53,8 +52,7 @@ def _iter_all_pairs(
         if not emb_path.exists():
             raise ValueError(f"Embedding file {emb_path} does not exist")
 
-        for record, emb in _pair_stream(txt_path, emb_path):
-            yield record, emb
+        yield from _pair_stream(txt_path, emb_path)
 
 
 def build_parquet_shards_from_folder(
@@ -73,8 +71,8 @@ def build_parquet_shards_from_folder(
     pair_iter = _iter_all_pairs(path)
     try:
         first_record, first_emb = next(pair_iter)
-    except StopIteration:
-        raise RuntimeError("No data found under the given path.")
+    except StopIteration as e:
+        raise RuntimeError("No data found under the given path.") from e
 
     keys = list(first_record.keys())
     first_key = keys[0]
@@ -168,8 +166,7 @@ def _post_process_dataset(dataset: Dataset | IterableDataset, to_keep: set[str])
 
 
 def _collect_parquet_shards(path_or_repo: Path) -> list[Path]:
-    """
-    Return a sorted list of train/*.parquet for local dir or HF dataset repo.
+    """Return a sorted list of train/*.parquet for local dir or HF dataset repo.
 
     The sort order is determined lexicographically by path.
     If a HF dataset repo is given, we first download a local snapshot, and then return any shards
@@ -186,7 +183,8 @@ def _collect_parquet_shards(path_or_repo: Path) -> list[Path]:
         shards = path_or_repo.glob("**/*.parquet")
     else:
         api = HfApi()  # pragma: no cover
-        local = Path(api.snapshot_download(repo_id=path_or_repo.as_posix(), repo_type="dataset"))  # pragma: no cover
+        path = api.snapshot_download(repo_id=path_or_repo.as_posix(), repo_type="dataset")
+        local = Path(path)  # pragma: no cover
         shards = local.glob("**/*.parquet")  # pragma: no cover
     return sorted(shards, key=lambda p: p.as_posix())
 
@@ -196,7 +194,7 @@ def get_datasets(
     paths: Sequence[str] | Sequence[Path],
     in_memory: Literal[True],
     limit_shards: int | None = None,
-    columns_to_keep: set[str] = {"sentence", "label", "question"},
+    columns_to_keep: frozenset[str] | set[str] = frozenset({"sentence", "label", "question"}),
 ) -> tuple[Dataset, int]: ...
 
 
@@ -205,7 +203,7 @@ def get_datasets(
     paths: Sequence[str] | Sequence[Path],
     in_memory: Literal[False],
     limit_shards: int | None = None,
-    columns_to_keep: set[str] = {"sentence", "label", "question"},
+    columns_to_keep: frozenset[str] | set[str] = frozenset({"sentence", "label", "question"}),
 ) -> tuple[IterableDataset, int]: ...
 
 
@@ -214,18 +212,17 @@ def get_datasets(
     paths: Sequence[str] | Sequence[Path],
     in_memory: bool,
     limit_shards: int | None = None,
-    columns_to_keep: set[str] = {"sentence", "label", "question"},
+    columns_to_keep: frozenset[str] | set[str] = frozenset({"sentence", "label", "question"}),
 ) -> tuple[IterableDataset | Dataset, int]: ...
 
 
 def get_datasets(
     paths: Sequence[Path] | Sequence[str],
     in_memory: bool = True,
     limit_shards: int | None = None,
-    columns_to_keep: set[str] = {"sentence", "label"},
+    columns_to_keep: frozenset[str] | set[str] = frozenset({"sentence", "label"}),
 ) -> tuple[Dataset | IterableDataset, int]:
-    """
-    Gets datasets from the given paths.
+    """Get datasets from the given paths.
 
     The datasets can be loaded in memory or streamed from disk. In either case, we assume that
     the datasets have a "train" split. In all cases, we assume that the datasets have "text" and "embedding"
@@ -270,7 +267,7 @@ def get_datasets(
         ds = cast(Dataset, ds)
         ds = ds.shuffle(seed=42)
 
-    dataset = _post_process_dataset(ds, to_keep=columns_to_keep)
+    dataset = _post_process_dataset(ds, to_keep=set(columns_to_keep))
     return dataset, length
 
 

pynife/dataset_vault.py

@@ -10,20 +10,19 @@ def _simple_text_field_dataset(
     config: str | None = None,
     split: str = "train",
 ) -> Iterator[dict[str, str]]:
-    """Helper function for datasets that extract a single text field."""
+    """Extract single text field."""
     dataset = cast(Dataset, load_dataset(huggingface_name, config, split=split))
 
     new_records: list[dict[str, str]] = []
     for record in cast(Iterable[dict[str, str]], dataset):
         text = record[text_field]
         new_records.append({"text": text})
 
-    return cast(Iterator[dict[str, str]], iter(new_records))
+    return iter(new_records)
 
 
 def english_words_definitions_dataset() -> tuple[str, str, Iterator[dict[str, str]]]:
-    """
-    Get the English Words Definitions dataset.
+    """Get the English Words Definitions dataset.
 
     Returns a tuple of (descriptive_name, huggingface_name, iterator).
     """
@@ -101,7 +100,7 @@ def snli_dataset() -> tuple[str, str, Iterator[dict[str, str]]]:
             seen.add(text)
             new_records.append({"text": text})
 
-    dataset_iterator = cast(Iterator[dict[str, str]], iter(new_records))
+    dataset_iterator = iter(new_records)
     return name, hf, dataset_iterator
 
 
@@ -216,8 +215,7 @@ def get_all_dataset_functions() -> dict[str, Callable[[], tuple[str, str, Iterat
 
 
 def short_dataset_name(hf_name: str) -> str:
-    """
-    Return the short dataset name from a HF hub identifier.
+    """Return the short dataset name from a HF hub identifier.
 
     Examples:
         - 'MongoDB/english-words-definitions' -> 'english-words-definitions'

pynife/distillation/infer.py

@@ -1,16 +1,16 @@
 import json
 import logging
 from collections.abc import Iterable, Iterator
+from contextlib import AbstractContextManager, nullcontext
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TypeVar
 
 import numpy as np
 import torch
 from sentence_transformers import SentenceTransformer
-from skeletoken import TokenizerModel
 from tqdm import tqdm
-from transformers import BatchEncoding, PreTrainedTokenizerBase, PreTrainedTokenizerFast
+from transformers import BatchEncoding, PreTrainedTokenizerBase
 
 from pynife.data import build_parquet_shards_from_folder
 from pynife.utilities import batchify
@@ -49,8 +49,7 @@ def _write_data(path: Path, pooled: list[torch.Tensor], records: list[dict[str,
 def _tokenize(
     strings: list[str], tokenizer: PreTrainedTokenizerBase, max_length: int
 ) -> tuple[BatchEncoding, list[str]]:
-    """
-    Tokenize a list of strings using a HuggingFace tokenizer.
+    """Tokenize a list of strings using a HuggingFace tokenizer.
 
     This is mainly a helper function; it also returns the truncated strings, so that we don't have to
     re-tokenize them later to find out how they were truncated.
@@ -79,22 +78,19 @@ def _tokenize(
     # Where the final dimension is start, end indices. So by taking the max end index
     # we know the length to which the tokenizer tokenized the string.
     lengths = np.asarray(offset_mapping)[:, :, 1].max(axis=1)
-    return tokenized, [string[:length] for string, length in zip(strings, lengths)]
+    return tokenized, [string[:length] for string, length in zip(strings, lengths, strict=False)]
 
 
 def _generate_embeddings(
     model: SentenceTransformer,
-    records: Iterator[dict[str, str]],
+    records: Iterator[dict[str, str]] | Iterable[dict[str, str]],
     output_dir: str | Path,
     batch_size: int = 96,
     max_length: int = 512,
     save_every: int = 8192,
     limit_batches: int | None = None,
-    lowercase: bool = True,
-    make_greedy: bool = True,
 ) -> None:
-    """
-    Generate embeddings for a stream of texts using a SentenceTransformer model.
+    """Generate embeddings for a stream of texts using a SentenceTransformer model.
 
     This is mainly used as an inner loop for the knowledge distillation process.
     We get N texts, and create embeddings for them using the teacher model.
@@ -110,8 +106,6 @@ def _generate_embeddings(
         max_length: The maximum sequence length for tokenization. Defaults to 512.
         save_every: Save intermediate results every N batches. Defaults to 8192.
         limit_batches: An optional limit on the number of batches to process.
-        lowercase: Whether to lowercase the tokenizer.
-        make_greedy: Whether to make the tokenizer greedy.
 
     """
     model.eval()
@@ -123,21 +117,13 @@ def _generate_embeddings(
     all_pooled, accumulated_records = [], []
     tokenizer: PreTrainedTokenizerBase = model.tokenizer
 
-    if lowercase and isinstance(tokenizer, PreTrainedTokenizerFast):
-        tokenizer_model = TokenizerModel.from_transformers_tokenizer(tokenizer)
-        tokenizer_model = tokenizer_model.decase_vocabulary()
-        tokenizer = tokenizer_model.to_transformers()
-    if make_greedy and isinstance(tokenizer, PreTrainedTokenizerFast):
-        tokenizer_model = TokenizerModel.from_transformers_tokenizer(tokenizer)
-        tokenizer_model = tokenizer_model.make_model_greedy()
-        tokenizer = tokenizer_model.to_transformers()
-
     original_max_length = model[0].max_seq_length
     assert original_max_length is not None
     assert isinstance(original_max_length, int)
     if max_length > original_max_length:
         logger.warning(
-            f"Warning: max_length {max_length} is greater than the model's max_length {original_max_length}. Not changing it."
+            f"Warning: max_length {max_length} is greater than the model's max_length {original_max_length}. "
+            "Not changing it."
         )
     else:
         model[0].max_seq_length = max_length  # type: ignore[assignment]
@@ -153,7 +139,7 @@ def _generate_embeddings(
             out = model(features_dict)
             pooled = out["sentence_embedding"].cpu()
 
-            for record, truncated in zip(batch, truncated_strings):
+            for record, truncated in zip(batch, truncated_strings, strict=False):
                 record["truncated"] = truncated
 
                 accumulated_records.append(record)
@@ -177,21 +163,32 @@ def _generate_embeddings(
         _write_data(path, all_pooled, accumulated_records, shards_saved)
 
 
+def _workdir(temp_folder: str | None) -> AbstractContextManager[str]:
+    """Create a temp folder if needed. For <= 3.11 temporary directories."""
+    if temp_folder is None:
+        return TemporaryDirectory()
+    Path(temp_folder).mkdir(parents=True, exist_ok=True)
+    return nullcontext(temp_folder)
+
+
 def generate_and_save_embeddings(
     model: SentenceTransformer,
     model_name: str,
     dataset_name: str,
     output_folder: str | Path,
-    records: Iterator[dict[str, str]],
+    records: Iterator[dict[str, str]] | Iterable[dict[str, str]],
+    temp_folder: str | None = None,
     limit_batches: int | None = None,
     batch_size: int = 512,
     save_every: int = 256,
     max_length: int = 512,
-    lowercase: bool = True,
-    make_greedy: bool = True,
 ) -> None:
-    """Run inference and save the results to parquet shards."""
-    with TemporaryDirectory() as dir_name:
+    """Run inference and save the results to parquet shards.
+
+    This runs inference over an iterable or iterator of records. Each record is a dictionary with
+    at the very least a "text" key, which is the field that will be featurized.
+    """
+    with _workdir(temp_folder) as dir_name:
         _generate_embeddings(
             model,
             records,
@@ -200,8 +197,6 @@ def generate_and_save_embeddings(
             save_every=save_every,
             limit_batches=limit_batches,
             max_length=max_length,
-            lowercase=lowercase,
-            make_greedy=make_greedy,
         )
 
         logger.info("Converting dataset to shards...")

pynife/embedding.py

@@ -61,7 +61,7 @@ def __init__(
         embedding_dim: int | None = None,
         **kwargs: Any,
     ) -> None:
-        """Static embedding layer."""
+        """Initialize static embedding layer."""
         super().__init__(tokenizer, embedding_weights, embedding_dim, **kwargs)
         self._max_seq_length = 512
 

pynife/losses.py

@@ -1,5 +1,5 @@
+from collections.abc import Sequence
 from enum import Enum
-from typing import Sequence
 
 import torch
 import torch.nn.functional as F
@@ -52,8 +52,10 @@ def select_loss(name: str | LossFunction) -> type[nn.Module]:
     """Select loss function by name."""
     try:
         function = LossFunction(name)
-    except ValueError:
-        raise ValueError(f"Unknown loss function: {name}, available options are: {[e.value for e in LossFunction]}")
+    except ValueError as e:
+        raise ValueError(
+            f"Unknown loss function: {name}, available options are: {[e.value for e in LossFunction]}"
+        ) from e
     match function:
         case LossFunction.COSINE:
             return CosineLoss

pynife/nife.py

@@ -5,8 +5,7 @@
 
 
 def load_as_router(name: str, teacher_name: str | None = None) -> SentenceTransformer:
-    """
-    Load a SentenceTransformer model from the Hugging Face Hub.
+    """Load a SentenceTransformer model from the Hugging Face Hub.
 
     Args:
         name: The name of the model to load.
@@ -33,5 +32,9 @@ def load_as_router(name: str, teacher_name: str | None = None) -> SentenceTransf
             "Please check that you have the correct teacher model."
         )
 
-    router = Router.for_query_document(query_modules=[small_model], document_modules=[big_model])  # type: ignore
+    router = Router.for_query_document(
+        query_modules=list(small_model),  # type: ignore  # BOOO
+        document_modules=list(big_model),  # type: ignore  # BOOO
+        default_route="query",
+    )
     return SentenceTransformer(modules=[router])