v0 litellm

mosaicml · Oct 5, 2023 · afde51d · afde51d
1 parent cb1d94a
commit afde51d
Showing 1 changed file with 4 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
@@ -81,14 +81,14 @@ class OpenAIChatAPIEvalWrapper(OpenAIEvalInterface):
     def __init__(self, model_cfg: Dict, tokenizer: AutoTokenizer) -> None:
         super().__init__(model_cfg, tokenizer)
         try:
-            import openai
+            import litellm
         except ImportError as e:
             raise MissingConditionalImportError(
-                extra_deps_group='openai',
-                conda_package='openai',
+                extra_deps_group='litellm',
+                conda_package='litellm',
                 conda_channel='conda-forge') from e
 
-        self.generate_completion = lambda prompt, num_tokens: openai.ChatCompletion.create(
+        self.generate_completion = lambda prompt, num_tokens: litellm.completion(
             self.model_name,
             messages=[{
                 'role': 'user',