Relax supported model check (#634)

Author: goodwanghan

.github/workflows/ci.yaml

@@ -40,7 +40,7 @@ jobs:
         run: python -c "from nixtla import NixtlaClient"
 
   run-all-tests:
-    runs-on: ubuntu-latest
+    runs-on: nixtla-linux-large-public
     timeout-minutes: 60
     strategy:
       fail-fast: false

nbs/src/nixtla_client.ipynb

@@ -46,6 +46,7 @@
     "import logging\n",
     "import math\n",
     "import os\n",
+    "import re\n",
     "import warnings\n",
     "from collections.abc import Sequence\n",
     "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
@@ -715,7 +716,34 @@
     "        if hist_exog_list and logger:\n",
     "            logger.info(f'Using historical exogenous features: {hist_exog_list}')\n",
     "\n",
-    "    return X, hist_exog"
+    "    return X, hist_exog\n",
+    "\n",
+    "def _model_in_list(model:str, model_list: tuple[Any]) -> bool:\n",
+    "    for m in model_list:\n",
+    "        if isinstance(m, str):\n",
+    "            if m == model:\n",
+    "                return True\n",
+    "        elif isinstance(m, re.Pattern):\n",
+    "            if m.fullmatch(model):\n",
+    "                return True\n",
+    "    return False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| hide\n",
+    "\n",
+    "assert _model_in_list(\"a\", (\"a\", \"b\"))\n",
+    "assert not _model_in_list(\"a\", (\"b\", \"c\"))\n",
+    "assert _model_in_list(\"axb\", (\"x\", re.compile(\"a.*b\")))\n",
+    "assert _model_in_list(\"axb\", (\"x\", re.compile(\"^a.*b$\")))\n",
+    "assert _model_in_list(\"a-b\", (\"x\", re.compile(\"^a-.*b$\")))\n",
+    "assert _model_in_list(\"a-dfdfb\", (\"x\", re.compile(\"^a-.*b$\")))\n",
+    "assert not _model_in_list(\"abc\", (\"x\", re.compile(\"ab\"), re.compile(\"abcd\")))"
    ]
   },
   {
@@ -1139,10 +1167,7 @@
     "        )\n",
     "        self._model_params: dict[tuple[str, str], tuple[int, int]] = {}\n",
     "        self._is_azure = 'ai.azure' in base_url\n",
-    "        if self._is_azure:\n",
-    "            self.supported_models = ['azureai']\n",
-    "        else:\n",
-    "            self.supported_models = ['timegpt-1', 'timegpt-1-long-horizon']\n",
+    "        self.supported_models:list[Any] = [re.compile('^timegpt-.+$'), 'azureai']\n",
     "\n",
     "    def _make_request(\n",
     "        self,\n",
@@ -1374,9 +1399,9 @@
     "    ) -> tuple[DFType, Optional[DFType], bool, _FreqType]:\n",
     "        if validate_api_key and not self.validate_api_key(log=False):\n",
     "            raise Exception('API Key not valid, please email support@nixtla.io')\n",
-    "        if model not in self.supported_models:\n",
+    "        if not _model_in_list(model, tuple(self.supported_models)):\n",
     "            raise ValueError(\n",
-    "                f'unsupported model: {model}. supported models: {self.supported_models}'\n",
+    "                f'unsupported model: {model}.'\n",
     "            )\n",
     "        drop_id = id_col not in df.columns\n",
     "        if drop_id:\n",

nixtla/__init__.py

@@ -1,3 +1,3 @@
-__version__ = "0.6.7.dev2"
+__version__ = "0.6.7.dev3"
 __all__ = ["NixtlaClient"]
 from .nixtla_client import NixtlaClient

nixtla/_modidx.py

@@ -129,6 +129,8 @@
                                                                                'nixtla/nixtla_client.py'),
                                       'nixtla.nixtla_client._maybe_infer_freq': ( 'src/nixtla_client.html#_maybe_infer_freq',
                                                                                   'nixtla/nixtla_client.py'),
+                                      'nixtla.nixtla_client._model_in_list': ( 'src/nixtla_client.html#_model_in_list',
+                                                                               'nixtla/nixtla_client.py'),
                                       'nixtla.nixtla_client._parse_in_sample_output': ( 'src/nixtla_client.html#_parse_in_sample_output',
                                                                                         'nixtla/nixtla_client.py'),
                                       'nixtla.nixtla_client._partition_series': ( 'src/nixtla_client.html#_partition_series',

nixtla/nixtla_client.py

@@ -9,6 +9,7 @@
 import logging
 import math
 import os
+import re
 import warnings
 from collections.abc import Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -638,15 +639,26 @@ def _process_exog_features(
 
     return X, hist_exog
 
-# %% ../nbs/src/nixtla_client.ipynb 11
+
+def _model_in_list(model: str, model_list: tuple[Any]) -> bool:
+    for m in model_list:
+        if isinstance(m, str):
+            if m == model:
+                return True
+        elif isinstance(m, re.Pattern):
+            if m.fullmatch(model):
+                return True
+    return False
+
+# %% ../nbs/src/nixtla_client.ipynb 12
 class AuditDataSeverity(Enum):
     """Enum class to indicate audit data severity levels"""
 
     FAIL = "Fail"  # Indicates a critical issue that requires immediate attention
     CASE_SPECIFIC = "Case Specific"  # Indicates an issue that may be acceptable in specific contexts
     PASS = "Pass"  # Indicates that the data is acceptable
 
-# %% ../nbs/src/nixtla_client.ipynb 13
+# %% ../nbs/src/nixtla_client.ipynb 14
 def _audit_duplicate_rows(
     df: AnyDFType,
     id_col: str = "unique_id",
@@ -660,7 +672,7 @@ def _audit_duplicate_rows(
     else:
         raise ValueError(f"Dataframe type {type(df)} is not supported yet.")
 
-# %% ../nbs/src/nixtla_client.ipynb 16
+# %% ../nbs/src/nixtla_client.ipynb 17
 def _audit_missing_dates(
     df: AnyDFType,
     freq: _Freq,
@@ -688,7 +700,7 @@ def _audit_missing_dates(
     else:
         raise ValueError(f"Dataframe type {type(df)} is not supported yet.")
 
-# %% ../nbs/src/nixtla_client.ipynb 19
+# %% ../nbs/src/nixtla_client.ipynb 20
 def _audit_categorical_variables(
     df: AnyDFType,
     id_col: str = "unique_id",
@@ -708,7 +720,7 @@ def _audit_categorical_variables(
     else:
         raise ValueError(f"Dataframe type {type(df)} is not supported yet.")
 
-# %% ../nbs/src/nixtla_client.ipynb 22
+# %% ../nbs/src/nixtla_client.ipynb 23
 def _audit_leading_zeros(
     df: pd.DataFrame,
     id_col: str = "unique_id",
@@ -733,7 +745,7 @@ def _audit_leading_zeros(
     else:
         raise ValueError(f"Dataframe type {type(df)} is not supported yet.")
 
-# %% ../nbs/src/nixtla_client.ipynb 25
+# %% ../nbs/src/nixtla_client.ipynb 26
 def _audit_negative_values(
     df: AnyDFType,
     target_col: str = "y",
@@ -746,7 +758,7 @@ def _audit_negative_values(
     else:
         raise ValueError(f"Dataframe type {type(df)} is not supported yet.")
 
-# %% ../nbs/src/nixtla_client.ipynb 27
+# %% ../nbs/src/nixtla_client.ipynb 28
 class ApiError(Exception):
     status_code: Optional[int]
     body: Any
@@ -760,7 +772,7 @@ def __init__(
     def __str__(self) -> str:
         return f"status_code: {self.status_code}, body: {self.body}"
 
-# %% ../nbs/src/nixtla_client.ipynb 29
+# %% ../nbs/src/nixtla_client.ipynb 30
 class NixtlaClient:
 
     def __init__(
@@ -821,10 +833,7 @@ def __init__(
         )
         self._model_params: dict[tuple[str, str], tuple[int, int]] = {}
         self._is_azure = "ai.azure" in base_url
-        if self._is_azure:
-            self.supported_models = ["azureai"]
-        else:
-            self.supported_models = ["timegpt-1", "timegpt-1-long-horizon"]
+        self.supported_models: list[Any] = [re.compile("^timegpt-.+$"), "azureai"]
 
     def _make_request(
         self,
@@ -1057,10 +1066,8 @@ def _run_validations(
     ) -> tuple[DFType, Optional[DFType], bool, _FreqType]:
         if validate_api_key and not self.validate_api_key(log=False):
             raise Exception("API Key not valid, please email support@nixtla.io")
-        if model not in self.supported_models:
-            raise ValueError(
-                f"unsupported model: {model}. supported models: {self.supported_models}"
-            )
+        if not _model_in_list(model, tuple(self.supported_models)):
+            raise ValueError(f"unsupported model: {model}.")
         drop_id = id_col not in df.columns
         if drop_id:
             df = ufp.copy_if_pandas(df, deep=False)
@@ -2862,7 +2869,7 @@ def clean_data(
 
         return df, all_pass, error_dfs, case_specific_dfs
 
-# %% ../nbs/src/nixtla_client.ipynb 31
+# %% ../nbs/src/nixtla_client.ipynb 32
 def _forecast_wrapper(
     df: pd.DataFrame,
     client: NixtlaClient,

settings.ini

@@ -8,7 +8,7 @@ author = Nixtla
 author_email = business@nixtla.io
 copyright = Nixtla Inc.
 branch = main
-version = 0.6.7.dev2
+version = 0.6.7.dev3
 min_python = 3.9
 audience = Developers
 language = English