chatbot-app / common /utility /embedding_factory.py
mishrabp's picture
Upload folder using huggingface_hub
6b50ab8 verified
import os
from typing import Union
# from azure.identity import DefaultAzureCredential
from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
from langchain_ollama import OllamaEmbeddings
from langchain_huggingface import HuggingFaceEmbeddings
class EmbeddingFactory:
"""
A static utility class to create and return LLM Embedding instances based on the input type.
"""
@staticmethod
def get_llm(llm_type: str) -> Union[AzureOpenAIEmbeddings, OpenAIEmbeddings]:
"""
Returns an LLM instance based on the specified type.
Parameters:
llm_type (str): The type of LLM to return. Valid values are 'azure' or 'openai'.
Returns:
Union[AzureOpenAIEmbeddings, OpenAIEmbeddings]: The LLM instance.
"""
if llm_type.lower() == "azure":
# Get the Azure Credential
# credential = DefaultAzureCredential()
# token=credential.get_token("https://cognitiveservices.azure.com/.default").token
# if not token:
# raise ValueError("Token is required for AzureOpenAIEmbeddings.")
# return AzureOpenAIEmbeddings(
# azure_endpoint=os.environ["AZURE_OPENAI_API_URI"],
# azure_deployment="text-embedding-3-small", #os.environ["AZURE_OPENAI_API_BASE_MODEL"],
# api_version=os.environ["AZURE_OPENAI_API_VERSION"],
# api_key=token
# )
pass
elif llm_type.lower() == "openai":
return OpenAIEmbeddings(
api_key=os.environ["OPENAI_API_KEY"],
model="text-embedding-3-large"
)
elif llm_type.lower() == "ollama": # must have ollama running locally with the following model
return OllamaEmbeddings(model="gemma:2b")
elif llm_type.lower() == "hf": # must have key update in env:HF_TOKEN
return HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
else:
raise ValueError("Invalid llm_type. Use 'azure' or 'openai'.")