fix: add skeletoken

Author: stephantul

model2vec/tokenizer/tokenizer.py

@@ -56,18 +56,7 @@ def replace_vocabulary(tokenizer: Tokenizer, new_vocabulary: list[Token]) -> Tok
         new_added_tokens.append(added_token)
     for token in new_vocabulary:
         if token.is_multiword and token.form not in {tokenizer_model.unk_token, tokenizer_model.pad_token}:
-            token_id = tokenizer_model.model.vocab[token.form]
-            new_added_tokens.append(
-                AddedToken(
-                    content=token.form,
-                    single_word=False,
-                    lstrip=True,
-                    rstrip=True,
-                    normalized=True,
-                    special=False,
-                    id=token_id,
-                )
-            )
+            tokenizer_model.add_addedtoken(token.form, normalized=True, single_word=False)
 
     pre_tokenized_tokens = [x.normalized_form for x in new_vocabulary]
     tokenizer_model.added_tokens = AddedTokens(_remap_added_tokens(new_added_tokens, pre_tokenized_tokens))
@@ -234,32 +223,32 @@ def turn_tokens_into_ids(tokens: list[Token], tokenizer: Tokenizer) -> list[list
     """
     prefix, suffix = find_eos_bos(tokenizer)
 
-    prefix_id, suffix_id = None, None
+    prefix_ids, suffix_ids = None, None
     vocab = tokenizer.get_vocab()
     if prefix is not None:
-        prefix_id = vocab[prefix]
+        prefix_ids = [vocab[token] for token in prefix]
     if suffix is not None:
-        suffix_id = vocab[suffix]
+        suffix_ids = [vocab[token] for token in suffix]
 
     token_ids: list[list[int]] = []
     for token in tokens:
         token_sequence = []
-        if prefix_id is not None:
-            token_sequence.append(prefix_id)
+        if prefix_ids is not None:
+            token_sequence.extend(prefix_ids)
         if token.is_internal:
             token_id = vocab[token.form]
             token_sequence.append(token_id)
         else:
             token_sequence.extend(tokenizer.encode(token.form).ids)
-        if suffix_id is not None:
-            token_sequence.append(suffix_id)
+        if suffix_ids is not None:
+            token_sequence.extend(suffix_ids)
 
         token_ids.append(token_sequence)
 
     return token_ids
 
 
-def find_eos_bos(tokenizer: Tokenizer) -> tuple[str | None, str | None]:
+def find_eos_bos(tokenizer: Tokenizer) -> tuple[list[str] | None, list[str] | None]:
     """Finds the eos and bos tokens for a tokenizer."""
     model = TokenizerModel.from_tokenizer(tokenizer)
     return model.bos, model.eos

model2vec/utils.py

@@ -1,18 +1,14 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-import json
 import logging
 import re
 from importlib import import_module
 from importlib.metadata import metadata
-from pathlib import Path
 from typing import Any, Iterator, Protocol, cast
 
 import numpy as np
-import safetensors
 from joblib import Parallel
-from tokenizers import Tokenizer
 from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
@@ -78,12 +74,15 @@ def get_package_extras(package: str, extra: str) -> Iterator[str]:
             found_extra = rest[0].split("==")[-1].strip(" \"'")
             if found_extra == extra:
                 prefix, *_ = _DIVIDERS.split(name)
+                prefix = prefix.split("@")[0].strip()
                 yield prefix.strip()
 
 
 def importable(module: str, extra: str) -> None:
     """Check if a module is importable."""
     module = dict(_MODULE_MAP).get(module, module)
+    # Allows this to work with git installed modules.
+    module = module.split("@")[0].strip()
     try:
         import_module(module)
     except ImportError:

pyproject.toml

@@ -60,7 +60,7 @@ dev = [
     "ruff",
 ]
 
-distill = ["torch", "transformers", "scikit-learn", "skeletoken @ git+https://github.com/stephantul/skeletoken.git"]
+distill = ["torch", "transformers", "scikit-learn", "skeletoken"]
 onnx = ["onnx", "torch"]
 # train also installs inference
 train = ["torch", "lightning", "scikit-learn", "skops"]

tests/test_utils.py

@@ -72,7 +72,7 @@ def test_importable() -> None:
 def test_get_package_extras() -> None:
     """Test package extras."""
     extras = set(get_package_extras("model2vec", "distill"))
-    assert extras == {"torch", "transformers", "scikit-learn"}
+    assert extras == {"torch", "transformers", "scikit-learn", "skeletoken"}
 
 
 def test_get_package_extras_empty() -> None: