fix bug

bmtools/agent/translator.py

@@ -2,7 +2,7 @@ from langchain.llms import OpenAI
 from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
 
-from bmtools.models import CustomLLM
+# from bmtools.models import CustomLLM
 
 import py3langid as langid
 from iso639 import languages
@@ -41,12 +41,12 @@ class Translator:
             return inputs
 
     def create_openai_model(self, openai_api_key: str, model_name: str) -> OpenAI:
-        # if openai_api_key is None:
-        #     openai_api_key = os.environ.get('OPENAI_API_KEY')
-        # llm = OpenAI(model_name=model_name,
-        #              temperature=0.0,
-        #              openai_api_key=openai_api_key)
-        llm = CustomLLM()
+        if openai_api_key is None:
+            openai_api_key = os.environ.get('OPENAI_API_KEY')
+        llm = OpenAI(model_name=model_name,
+                     temperature=0.0,
+                     openai_api_key=openai_api_key)
+        # llm = CustomLLM()
         return llm
 
     def create_prompt(self) -> PromptTemplate: