mirror of
https://github.com/qdrant/fastembed.git
synced 2026-10-02 19:07:53 -05:00
Andrey VasnetsovandGeorge
04bc7a3039
* add token embeddings * fix parallel worker init * implement minicoil * fix mypy * fix mypy * register minicoil * rollback suggested "fix" * add minicoil test * some pr issues (#514) * some pr issues * revert query embed refactor * test: add query embed tests * nit * Update tests/test_sparse_embeddings.py --------- Co-authored-by: Andrey Vasnetsov <andrey@vasnetsov.com> * review * fix: revert change to colbert query_embed --------- Co-authored-by: George <george.panchuk@qdrant.tech>
84 lines
3.2 KiB
Python
84 lines
3.2 KiB
Python
import json
|
|
from typing import Any
|
|
from pathlib import Path
|
|
|
|
from tokenizers import AddedToken, Tokenizer
|
|
|
|
from fastembed.image.transform.operators import Compose
|
|
|
|
|
|
def load_special_tokens(model_dir: Path) -> dict[str, Any]:
|
|
tokens_map_path = model_dir / "special_tokens_map.json"
|
|
if not tokens_map_path.exists():
|
|
raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
|
|
|
|
with open(str(tokens_map_path)) as tokens_map_file:
|
|
tokens_map = json.load(tokens_map_file)
|
|
|
|
return tokens_map
|
|
|
|
|
|
def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:
|
|
config_path = model_dir / "config.json"
|
|
if not config_path.exists():
|
|
raise ValueError(f"Could not find config.json in {model_dir}")
|
|
|
|
tokenizer_path = model_dir / "tokenizer.json"
|
|
if not tokenizer_path.exists():
|
|
raise ValueError(f"Could not find tokenizer.json in {model_dir}")
|
|
|
|
tokenizer_config_path = model_dir / "tokenizer_config.json"
|
|
if not tokenizer_config_path.exists():
|
|
raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
|
|
|
|
with open(str(config_path)) as config_file:
|
|
config = json.load(config_file)
|
|
|
|
with open(str(tokenizer_config_path)) as tokenizer_config_file:
|
|
tokenizer_config = json.load(tokenizer_config_file)
|
|
assert "model_max_length" in tokenizer_config or "max_length" in tokenizer_config, (
|
|
"Models without model_max_length or max_length are not supported."
|
|
)
|
|
if "model_max_length" not in tokenizer_config:
|
|
max_context = tokenizer_config["max_length"]
|
|
elif "max_length" not in tokenizer_config:
|
|
max_context = tokenizer_config["model_max_length"]
|
|
else:
|
|
max_context = min(tokenizer_config["model_max_length"], tokenizer_config["max_length"])
|
|
|
|
tokens_map = load_special_tokens(model_dir)
|
|
|
|
tokenizer = Tokenizer.from_file(str(tokenizer_path))
|
|
tokenizer.enable_truncation(max_length=max_context)
|
|
tokenizer.enable_padding(
|
|
pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
|
|
)
|
|
|
|
for token in tokens_map.values():
|
|
if isinstance(token, str):
|
|
tokenizer.add_special_tokens([token])
|
|
elif isinstance(token, dict):
|
|
tokenizer.add_special_tokens([AddedToken(**token)])
|
|
|
|
special_token_to_id: dict[str, int] = {}
|
|
|
|
for token in tokens_map.values():
|
|
if isinstance(token, str):
|
|
special_token_to_id[token] = tokenizer.token_to_id(token)
|
|
elif isinstance(token, dict):
|
|
token_str = token.get("content", "")
|
|
special_token_to_id[token_str] = tokenizer.token_to_id(token_str)
|
|
|
|
return tokenizer, special_token_to_id
|
|
|
|
|
|
def load_preprocessor(model_dir: Path) -> Compose:
|
|
preprocessor_config_path = model_dir / "preprocessor_config.json"
|
|
if not preprocessor_config_path.exists():
|
|
raise ValueError(f"Could not find preprocessor_config.json in {model_dir}")
|
|
|
|
with open(str(preprocessor_config_path)) as preprocessor_config_file:
|
|
preprocessor_config = json.load(preprocessor_config_file)
|
|
transforms = Compose.from_config(preprocessor_config)
|
|
return transforms
|