File size: 502 Bytes
b0bfa1a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
import os
from langchain import OpenAI
from langchain.chat_models import ChatOpenAI
def get_llm(llm_name, model_temperature, api_key, max_tokens=256):
os.environ["OPENAI_API_KEY"] = api_key
if llm_name == "text-davinci-003":
return OpenAI(
temperature=model_temperature, model_name=llm_name, max_tokens=max_tokens
)
else:
return ChatOpenAI(
temperature=model_temperature, model_name=llm_name, max_tokens=max_tokens
)
|