1from functools import partial
2
3import numpy as np
4from llama_cpp import Llama
5
6max_length = 512
7
8model = Llama.from_pretrained(
9 repo_id="mm/japanese-e5-mistral-7b_slerp_gguf",
10 filename="*Q4_K.gguf", # Choose from the avaiable formats,
11 embedding=True,
12 n_ctx=max_length,
13 n_batch=max_length,
14 verbose=False,
15)
16model.tokenize = partial(model.tokenize, special=True)
17
18
19def calc_emb(s: str):
20 if len(model.tokenize(s.encode())) > max_length - 1:
21 print(
22 "The output will be calculated with truncation because of the length exceeding."
23 )
24 v = model.embed(s + "</s>", normalize=True, truncate=True)
25 return np.asarray(v[-1])
26
27
28s = "今日の天気は？"
29t = "本日の天候は？"
30
31print(f"cossim({s}, {t}) = {(calc_emb(s) * calc_emb(t)).sum()}")
