File size: 2,094 Bytes
d9674cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | from pathlib import Path
import torch
from sentence_transformers import SentenceTransformer
BASE_MODEL_NAME = "BAAI/bge-m3"
CHECKPOINT_FILENAME = "bge_m3_quantized_state_dict.pt"
MAX_SEQ_LENGTH = 128
def _quantize_model(model):
try:
from torch.ao.quantization import quantize_dynamic
except ImportError:
from torch.quantization import quantize_dynamic
result = quantize_dynamic(
model,
{torch.nn.Linear},
dtype=torch.qint8,
inplace=True,
)
return result if result is not None else model
def load_model(checkpoint_path=None):
if checkpoint_path is None:
checkpoint_path = (
Path(__file__).resolve().parent
/ CHECKPOINT_FILENAME
)
else:
checkpoint_path = Path(checkpoint_path)
if not checkpoint_path.is_file():
raise FileNotFoundError(
f"Checkpoint не найден: {checkpoint_path}"
)
model = SentenceTransformer(
BASE_MODEL_NAME,
device="cpu",
)
model.max_seq_length = MAX_SEQ_LENGTH
model.eval()
model = _quantize_model(model)
model.max_seq_length = MAX_SEQ_LENGTH
model.eval()
try:
state_dict = torch.load(
checkpoint_path,
map_location="cpu",
weights_only=False,
)
except TypeError:
state_dict = torch.load(
checkpoint_path,
map_location="cpu",
)
model.load_state_dict(
state_dict,
strict=True,
)
model.eval()
return model
if __name__ == "__main__":
model = load_model()
embeddings = model.encode(
["Первый текст", "Второй текст"],
batch_size=2,
show_progress_bar=False,
normalize_embeddings=True,
convert_to_numpy=True,
device="cpu",
)
print("Модель успешно загружена.")
print("Размер эмбеддингов:", embeddings.shape)
|