File size: 905 Bytes
1ff6afc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
import yaml
from src.llm.enums import OPENAI_LLM, NVIDIA_LLM
from src.llm.base_llm_provider import BaseLLMProvider
from src.llm.openai_llm import OpenAILLM
from src.llm.nvidia_llm import NvidiaLLM
def get_llm(config_file_path: str = "config.yaml") -> BaseLLMProvider:
"""
Initiates LLM client from config file
"""
# load config
with open(config_file_path, "r") as f:
config = yaml.safe_load(f)
# init & return llm
if config["PROVIDER"] == OPENAI_LLM:
return OpenAILLM(
model=config["MODEL"],
temperature=config["TEMPERATURE"],
base_url=config["BASE_URL"],
)
elif config["PROVIDER"] == NVIDIA_LLM:
return NvidiaLLM(
model=config["MODEL"],
temperature=config["TEMPERATURE"],
base_url=config["BASE_URL"],
)
else:
raise ValueError(config["MODEL"])
|