src/embeddings.py

# Databricks notebook source
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores.utils import DistanceStrategy


def init_embedding_model(EMBEDDING_MODEL_NAME: str):
    embedding_model = HuggingFaceEmbeddings(
        model_name=EMBEDDING_MODEL_NAME,
        multi_process=True,
        model_kwargs={"device": "cuda"},
        # model_kwargs={"device": "cpu"},
        # Set `True` for cosine similarity
        encode_kwargs={"normalize_embeddings": True},
    )
    return embedding_model