diff --git a/llmfoundry/models/inference_api_wrapper/openai_causal_lm.py b/llmfoundry/models/inference_api_wrapper/openai_causal_lm.py index 7257b98bd8..311f35d3f3 100644 --- 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 index 83fecaff31..a43925e506 100644 --- 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