[py] test Ruri text embeddings

1vectorize.py

import argparse

parser = argparse.ArgumentParser(description='Process text file and create tensor embeddings')
parser.add_argument('textfile', help='Input text file path')
parser.add_argument('--ollama', action='store_true', help='Use Ollama')
arg = parser.parse_args()

import os, torch, safetensors.torch
from tqdm import tqdm

if arg.ollama:
    import ollama
    model = "kun432/cl-nagoya-ruri-large"
    def embed(s):
        return torch.tensor([ollama.embeddings(model=model, prompt=s).embedding])
else:
    from sentence_transformers import SentenceTransformer
    # Download from the 🤗 Hub
    model = SentenceTransformer("cl-nagoya/ruri-base")
    def embed(s):
        return model.encode([s], convert_to_tensor=True)

tensorfile = os.path.splitext(arg.textfile)[0] + ".safetensors"

# Don't forget to add the prefix "クエリ: " for query-side or "文章: " for passage-side texts.
with open(arg.textfile, "r", encoding="utf-8") as f:
    lines = [l for line in f if (l := line.strip())]
test = embed("文章: test")[0]
print("vector size:", len(test))
tensor = torch.zeros(len(lines), len(test), dtype=torch.float32)
for i, line in tqdm(enumerate(lines), total=len(lines)):
    # print(f"{i+1} / {len(lines)} {line}")
    tensor[i, :] = embed(f"文章: {line}")[0]
safetensors.torch.save_file({"lines": tensor}, tensorfile)

2test.py

import argparse

parser = argparse.ArgumentParser(description='Process text file and create tensor embeddings')
parser.add_argument('textfile', help='Input text file path')
parser.add_argument('--ollama', action='store_true', help='Use Ollama')
arg = parser.parse_args()

import os, torch, torch.nn.functional as F, safetensors.torch

if arg.ollama:
    import ollama
    model = "kun432/cl-nagoya-ruri-large"
    def embed(s):
        return torch.tensor([ollama.embeddings(model=model, prompt=s).embedding])
else:
    from sentence_transformers import SentenceTransformer
    # Download from the 🤗 Hub
    model = SentenceTransformer("cl-nagoya/ruri-base")
    def embed(s):
        return model.encode([s], convert_to_tensor=True)

tensorfile = os.path.splitext(arg.textfile)[0] + ".safetensors"

with open(arg.textfile, "r", encoding="utf-8") as f:
    lines = [l for line in f if (l := line.strip())]
tensor = safetensors.torch.load_file(tensorfile)["lines"]

# Don't forget to add the prefix "クエリ: " for query-side or "文章: " for passage-side texts.
while True:
    print()
    try:
        q = input("> ")
    except:
        print()
        break
    embeddings = embed(f"クエリ: {q}")
    similarities = F.cosine_similarity(tensor, embeddings, dim=1)
    for i, (value, index) in enumerate(zip(*torch.topk(similarities, k=10))):
        v, idx = value.item(), index.item()
        print(f"{i+1:2d}: {v:.5f} {idx + 1:4d} {lines[idx]}")