src/apps/utils/embeddings.py

import os
from transformers import AutoModel
from numpy.linalg import norm

cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b))
_model = None

def get_model():
    global _model
    if _model is None:
        BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
        model_path = os.path.join(BASE_DIR, "volumes", "models", "jina-embeddings-v2-base-en")
        
        if os.path.exists(model_path):
            print(f"Loading the model weights from local path: {model_path}")
            _model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
        else:
            print("Local model weights not found. Downloading from Hugging Face Hub (jinaai/jina-embeddings-v2-base-en)...")
            _model = AutoModel.from_pretrained("jinaai/jina-embeddings-v2-base-en", trust_remote_code=True)
    return _model

def get_embeddings(text:list):
    model = get_model()
    embeddings = model.encode(text)
    normalized_embeddings = embeddings/norm(embeddings[0])
    return normalized_embeddings