transforms/embed.py

import os

from typing import Union

from tclogger import logger
from transformers import AutoModel
from numpy.linalg import norm

from configs.envs import ENVS

os.environ["HF_ENDPOINT"] = ENVS["HF_ENDPOINT"]
os.environ["HF_TOKEN"] = ENVS["HF_TOKEN"]


def cosine_similarity(a, b):
    return (a @ b.T) / (norm(a) * norm(b))


class JinaAIEmbedder:
    def __init__(self, model_name: str = "jinaai/jina-embeddings-v2-base-zh"):
        self.model_name = model_name
        self.load_model()

    def load_model(self):
        self.model = AutoModel.from_pretrained(self.model_name, trust_remote_code=True)

    def encode(self, text: Union[str, list[str]]):
        if isinstance(text, str):
            text = [text]
        return self.model.encode(text)


if __name__ == "__main__":
    embedder = JinaAIEmbedder()
    text = ["How is the weather today?", "今天天气怎么样?"]
    # text = "How is the weather today?"
    embeddings = embedder.encode(text)
    logger.success(embeddings)
    # print(cosine_similarity(embeddings[0], embeddings[1]))