get_embedding_function.py

from langchain_ollama import OllamaEmbeddings
from langchain_community.embeddings.bedrock import BedrockEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings
from sentence_transformers import SentenceTransformer
import time
import numpy as np
import logging
from FlagEmbedding import FlagModel


def get_embedding_function(code):

    if(code == "all-MiniLM-L6-v2"):
        model_name = "all-MiniLM-L6-v2"
        # You might need to specify encode_kwargs if normalizing embeddings
        # encode_kwargs = {'normalize_embeddings': True} # Example
        embeddings = HuggingFaceEmbeddings(
            model_name=model_name,
            model_kwargs={'device': 'cuda'} # Or 'cpu' - specify device if needed
            # encode_kwargs=encode_kwargs # Add if needed
        )
        return embeddings
    elif(code == "all-mpnet-base-v2"):
        model_name = "all-mpnet-base-v2"
        embeddings = HuggingFaceEmbeddings(
            model_name=model_name,
            model_kwargs={'device': 'cuda'} # Or 'cpu' - specify device if needed
        )
        return embeddings
    elif(code == "bge/"):
        model = FlagModel('BAAI/bge-small-en-v1.5', query_instruction_for_retrieval="Represent this sentence for searching relevant passages:")
        embeddings = model.encode
        return embeddings 
    elif(code == "ollama"):
        embeddings = OllamaEmbeddings(model="mistral", temperature=0.1)
        return embeddings