add it

mosaicml · Dec 15, 2023 · c01f9d9 · c01f9d9
1 parent cc60d01
commit c01f9d9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/llmfoundry/models/inference_api_wrapper/openai_causal_lm.py b/llmfoundry/models/inference_api_wrapper/openai_causal_lm.py
@@ -23,9 +23,11 @@
     'OpenAICausalLMEvalWrapper',
     'OpenAIChatAPIEvalWrapper',
 ]
-from openai.types.chat.chat_completion import ChatCompletion
-from openai.types.completion import Completion
-from openai.types.completion_choice import Logprobs
+
+if TYPE_CHECKING:
+    from openai.types.chat.chat_completion import ChatCompletion
+    from openai.types.completion import Completion
+    from openai.types.completion_choice import Logprobs
 
 MAX_RETRIES = 10
 

diff --git a/tests/test_smoketest.py b/tests/test_smoketest.py
@@ -1,7 +1,7 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
-from llmfoundry import (callbacks, data, models, optim, tokenizers, utils)
+from llmfoundry import callbacks, data, models, optim, tokenizers, utils
 
 
 # This very simple test is just to use the above imports, which check and make sure we can import all the top-level