Make numpy and pandas optional for ~7 times smaller deps

openai/api_resources/embedding.py

@@ -1,11 +1,10 @@
 import base64
 import time
 
-import numpy as np
 
 from openai import util
-from openai.api_resources.abstract import DeletableAPIResource, ListableAPIResource
 from openai.api_resources.abstract.engine_api_resource import EngineAPIResource
+from openai.datalib import numpy as np, assert_has_numpy
 from openai.error import TryAgain
 
 
@@ -40,6 +39,7 @@ def create(cls, *args, **kwargs):
 
                         # If an engine isn't using this optimization, don't do anything
                         if type(data["embedding"]) == str:
+                            assert_has_numpy()
                             data["embedding"] = np.frombuffer(
                                 base64.b64decode(data["embedding"]), dtype="float32"
                             ).tolist()

openai/datalib.py

@@ -0,0 +1,40 @@
+"""
+This module helps make data libraries like `numpy` and `pandas` optional dependencies.
+
+The libraries add up to 130MB+, which makes it challenging to deploy applications
+using this library in environments with code size constraints, like AWS Lambda.
+
+This module serves as an import proxy and provides a few utilities for dealing with the optionality.
+
+Since the primary use case of this library (talking to the OpenAI API) doesn’t generally require data libraries,
+it’s safe to make them optional. The rare case when data libraries are needed in the client is handled through
+assertions with instructive error messages.
+
+See also `setup.py`.
+
+"""
+try:
+    import numpy
+except ImportError:
+    numpy = None
+
+try:
+    import pandas
+except ImportError:
+    pandas = None
+
+HAS_NUMPY = bool(numpy)
+HAS_PANDAS = bool(pandas)
+
+NUMPY_INSTRUCTIONS = "numpy is not installed: pip install openai[datalib]"
+PANDAS_INSTRUCTIONS = "pandas is not installed: pip install openai[datalib]"
+
+
+def assert_has_numpy():
+    if not HAS_NUMPY:
+        raise Exception(NUMPY_INSTRUCTIONS)
+
+
+def assert_has_pandas():
+    if not HAS_PANDAS:
+        raise Exception(PANDAS_INSTRUCTIONS)

openai/embeddings_utils.py

@@ -2,8 +2,6 @@
 from typing import List, Optional
 
 import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
 import plotly.express as px
 from scipy import spatial
 from sklearn.decomposition import PCA
@@ -12,6 +10,8 @@
 from tenacity import retry, stop_after_attempt, wait_random_exponential
 
 import openai
+from openai.datalib import numpy as np
+from openai.datalib import pandas as pd
 
 
 @retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))

openai/tests/test_long_examples_validator.py

@@ -2,9 +2,14 @@
 import subprocess
 from tempfile import NamedTemporaryFile
 
+import pytest
+
+from openai.datalib import HAS_PANDAS, HAS_NUMPY, NUMPY_INSTRUCTIONS, PANDAS_INSTRUCTIONS
 
-def test_long_examples_validator() -> None:
 
+@pytest.mark.skipif(not HAS_PANDAS, reason=PANDAS_INSTRUCTIONS)
+@pytest.mark.skipif(not HAS_NUMPY, reason=NUMPY_INSTRUCTIONS)
+def test_long_examples_validator() -> None:
     """
     Ensures that long_examples_validator() handles previously applied recommendations,
     namely dropped duplicates, without resulting in a KeyError.
@@ -45,4 +50,4 @@ def test_long_examples_validator() -> None:
     # validate get_long_indexes() applied during optional_fn() call in long_examples_validator()
     assert "indices of the long examples has changed" in prepared_data_cmd_output.stdout
 
-    return prepared_data_cmd_output.stdout
+    return prepared_data_cmd_output.stdout

openai/validators.py

@@ -2,7 +2,7 @@
 import sys
 from typing import Any, Callable, NamedTuple, Optional
 
-import pandas as pd
+from openai.datalib import pandas as pd, assert_has_pandas
 
 
 class Remediation(NamedTuple):
@@ -474,6 +474,7 @@ def read_any_format(fname, fields=["prompt", "completion"]):
      - for .xlsx it will read the first sheet
      - for .txt it will assume completions and split on newline
     """
+    assert_has_pandas()
     remediation = None
     necessary_msg = None
     immediate_msg = None

openai/wandb_logger.py

@@ -13,10 +13,9 @@
     import re
     from pathlib import Path
 
-    import numpy as np
-    import pandas as pd
-
     from openai import File, FineTune
+    from openai.datalib import numpy as np
+    from openai.datalib import pandas as pd
 
 
 class WandbLogger:

setup.py

@@ -9,28 +9,42 @@
 with open(version_path, "rt") as f:
     exec(f.read(), version_contents)
 
+
+DATA_LIBRARIES = [
+    # These libraries are optional because of their size. See `openai/datalib.py`.
+    "numpy",
+    "pandas>=1.2.3",  # Needed for CLI fine-tuning data preparation tool
+    "pandas-stubs>=1.1.0.11",  # Needed for type hints for mypy
+    "openpyxl>=3.0.7",  # Needed for CLI fine-tuning data preparation tool xlsx format
+]
+
 setup(
     name="openai",
     description="Python client library for the OpenAI API",
     version=version_contents["VERSION"],
     install_requires=[
         "requests>=2.20",  # to get the patch for CVE-2018-18074
         "tqdm",  # Needed for progress bars
-        "pandas>=1.2.3",  # Needed for CLI fine-tuning data preparation tool
-        "pandas-stubs>=1.1.0.11",  # Needed for type hints for mypy
-        "openpyxl>=3.0.7",  # Needed for CLI fine-tuning data preparation tool xlsx format
-        "numpy",
         "typing_extensions",  # Needed for type hints for mypy
     ],
     extras_require={
-        "dev": ["black~=21.6b0", "pytest==6.*"],
-        "wandb": ["wandb"],
+        "dev": [
+            "black~=21.6b0",
+            "pytest==6.*",
+            "pytest_mock",
+        ],
+        "datalib": DATA_LIBRARIES,
+        "wandb": [
+            "wandb",
+            *DATA_LIBRARIES,
+        ],
         "embeddings": [
             "scikit-learn>=1.0.2",  # Needed for embedding utils, versions >= 1.1 require python 3.8
             "tenacity>=8.0.1",
             "matplotlib",
             "sklearn",
             "plotly",
+            *DATA_LIBRARIES,
         ],
     },
     python_requires=">=3.7.1",