mirror of
https://github.com/qdrant/fastembed.git
synced 2026-07-23 11:20:51 -05:00
331207976e
* MiniLM fix * Added MiniLM to text embedding Fixed MiniLM source destination Black + isort for repo * Fixed model all-MiniLM-L6-v2 description Recomputed canonical vector for all-MiniLM-L6-v2 in test --------- Co-authored-by: d.rudenko <dimitriyrudenk@gmail.com>
77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
import json
|
|
from pathlib import Path
|
|
from typing import Tuple
|
|
|
|
from tokenizers import AddedToken, Tokenizer
|
|
|
|
from fastembed.image.transform.operators import Compose
|
|
|
|
|
|
def load_special_tokens(model_dir: Path) -> dict:
|
|
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, max_length: int = 512) -> Tuple[Tokenizer, dict]:
|
|
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)
|
|
|
|
tokens_map = load_special_tokens(model_dir)
|
|
|
|
tokenizer = Tokenizer.from_file(str(tokenizer_path))
|
|
tokenizer.enable_truncation(
|
|
max_length=min(tokenizer_config["model_max_length"], max_length)
|
|
)
|
|
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 = {}
|
|
|
|
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
|