Merge branch 'main' into vocquant

Author: stephantul

README.md

@@ -1,14 +1,10 @@
 
-<div align="center">
-    <picture>
-      <img width="35%" alt="Model2Vec logo" src="assets/images/logo_v2.png">
-    </picture>
-  </a>
-</div>
+<h2 align="center">
+  <img width="35%" alt="Model2Vec logo" src="assets/images/model2vec_logo.png"><br/>
+  Fast State-of-the-Art Static Embeddings
+</h2>
+
 
-<div align="center">
-  <h2>Fast State-of-the-Art Static Embeddings</h2>
-</div>
 
 <div align="center">
   <h2>

assets/images/model2vec_logo.png

@@ -371,7 +371,7 @@ def encode_as_sequence(
             return out_array[0]
         return out_array
 
-    def _encode_batch_as_sequence(self, sentences: list[str], max_length: int | None) -> list[np.ndarray]:
+    def _encode_batch_as_sequence(self, sentences: Sequence[str], max_length: int | None) -> list[np.ndarray]:
         """Encode a batch of sentences as a sequence."""
         ids = self.tokenize(sentences=sentences, max_length=max_length)
         out: list[np.ndarray] = []

model2vec/model.py

@@ -40,4 +40,4 @@ def _process_unigram(
 def _calculate_token_weight_for_unigram(token: str) -> float:
     """Calculate the token weight for Unigram."""
     # Always prefer longer tokens.
-    return len(token) + token.count("▁")
+    return len(token) + token.count("▁") + token.count("Ġ")

model2vec/tokenizer/model.py

@@ -65,7 +65,9 @@ def replace_vocabulary(
 
     # Remove old added tokens from added tokens
     tokenizer_json["added_tokens"] = [x for x in added_tokens if x["content"] in {"[UNK]", "[PAD]"}]
-    tokenizer_json = process_tokenizer(tokenizer_json, pre_tokenized_tokens, "[UNK]")
+    tokenizer_json = process_tokenizer(
+        tokenizer_json, pre_tokenized_tokens, "[UNK]" if "[UNK]" in pre_tokenized_tokens else None
+    )
 
     # Remap special tokens
     tokenizer_json["added_tokens"] = _remap_added_tokens(
@@ -111,11 +113,11 @@ def clean_and_create_vocabulary(
     internal_vocab: dict[str, int] = tokenizer.get_vocab()
     internal_tokens: list[str] = [k for k, _ in sorted(internal_vocab.items(), key=lambda x: x[1])]
 
+    cleaned_vocabulary = _process_internal_tokens(tokenizer, backend_tokenizer, internal_tokens, token_remove_regex)
     # Copy the backend tokenizer to avoid modifying the original.
     backend_tokenizer = backend_tokenizer.from_str(backend_tokenizer.to_str())
     backend_tokenizer = replace_normalizer(backend_tokenizer)
 
-    cleaned_vocabulary = _process_internal_tokens(tokenizer, backend_tokenizer, internal_tokens, token_remove_regex)
     internal_tokens_set = {token.form for token in cleaned_vocabulary}
 
     normalizer: Normalizer | None = backend_tokenizer.normalizer
@@ -302,15 +304,9 @@ def turn_tokens_into_ids(
     :param tokenizer: The tokenizer to use for converting tokens to IDs
     :param unk_token: The string form of the unk token.
     :return: List of token IDs corresponding to the input tokens
-    :raises ValueError: If the tokenizer returns an unexpected number of tokens for a single token
     """
     unk_id = None if unk_token is None else tokenizer.convert_tokens_to_ids(unk_token)
-
-    encoding = tokenizer.encode("a", add_special_tokens=True)
-
-    if len(encoding) != 3:
-        raise ValueError(f"Tokenizer returned {len(encoding)} tokens for a single token. This is not supported.")
-    bos, _, eos = encoding
+    prefix, suffix = find_eos_bos(tokenizer)
 
     token_ids: list[list[int]] = []
     for token in tokens:
@@ -321,13 +317,30 @@ def turn_tokens_into_ids(
             # Explicitly check and warn if `unk_id` appears, but don't crash.
             if unk_id is not None and token_id == unk_id and token.form != unk_token:
                 logger.warning(f"Token {token.form} was set to unk. This is wrong.")
-            token_ids.append([bos, token_id, eos])
+            token_ids.append([*prefix, token_id, *suffix])
         else:
             token_ids.append(tokenizer.encode(token.form))
 
     return token_ids
 
 
+def find_eos_bos(tokenizer: PreTrainedTokenizerFast) -> tuple[list[int], list[int]]:
+    """Finds the eos and bos tokens for a tokenizer."""
+    # Little bit complicated, because not all tokenizers have eos and bos tokens.
+    encoding = tokenizer.encode("a", add_special_tokens=True)
+    if len(encoding) != 3:
+        a_encoded = tokenizer.encode("a", add_special_tokens=False)
+        if len(a_encoded) != 1:
+            raise ValueError(
+                f"Error while encoding, couldn't determine eos and bos tokens. The model tokenizes 'a' to '{a_encoded}'"
+            )
+        a_idx = encoding.index(a_encoded[0])
+        prefix, suffix = encoding[:a_idx], encoding[a_idx + 1 :]
+    else:
+        prefix, suffix = encoding[:1], encoding[2:]
+    return prefix, suffix
+
+
 def _normalize_vocabulary_token(token: str, pre_tokenizer: PreTokenizer) -> str:
     """Normalize a token that is not in the initial token vocabulary."""
     # Add prefix space for byte tokenizers.