Fix some projects (#162)

* rm custom huggingface code

* fix role, syntax, and more

* upgrade deps

* use langchain community

* download erquired nltk data

* use langchain openai

* add faiss materializer

* failed attempt at agent executor materializer

* return agent and tools tuple

* update zenml version

* Update dependencies and README run command

* bump zenml version to the latest

* formatting

* formatting

* Fix NLTK data download to use a writable directory

* Retrieve OpenAI API key from ZenML secrets in index generator

* Reorder imports in web_url_loader.py

* Add environment variable fallback for OpenAI API key in index generator

* Add OpenAI API key retrieval in FAISS materializer

* Add OpenAI API key retrieval in agent creator

* update zenml syntax

* update code with right syntax and requirements

* update workspace

* update requirements

* update makefile and requirements

* make it use uv

* mount local cache dir to docker orchestrator

* add instructions for local_docker orch

---------

Co-authored-by: Alex Strick van Linschoten <[email protected]>

Author: wjayesh

airflow-cloud-composer-etl-feature-train/steps/etl/transform.py

@@ -18,10 +18,10 @@
 import os
 from datetime import datetime, timezone
 from typing import Optional
-from typing_extensions import Annotated
 
 import pandas as pd
 from materializers import BigQueryDataset, CSVDataset
+from typing_extensions import Annotated
 from zenml import step
 from zenml.logger import get_logger
 

airflow-cloud-composer-etl-feature-train/steps/feature_engineering/augment.py

@@ -18,9 +18,9 @@
 import os
 from datetime import datetime, timezone
 from typing import Optional
-from typing_extensions import Annotated
 
 from materializers import BigQueryDataset, CSVDataset
+from typing_extensions import Annotated
 from zenml import step
 from zenml.logger import get_logger
 

classifier-e2e/run.py

@@ -196,12 +196,12 @@ def main(
                 test_dataset_name, test_dataset_version_name
             )
             # Use versioned artifacts
-            run_args_train[
-                "train_dataset_id"
-            ] = train_dataset_artifact_version.id
-            run_args_train[
-                "test_dataset_id"
-            ] = test_dataset_artifact_version.id
+            run_args_train["train_dataset_id"] = (
+                train_dataset_artifact_version.id
+            )
+            run_args_train["test_dataset_id"] = (
+                test_dataset_artifact_version.id
+            )
 
         run_args_train["random_state"] = random.randint(0, 1000)
 

classifier-e2e/run_full.ipynb

@@ -41,6 +41,7 @@
     "! zenml login https://1cf18d95-zenml.cloudinfra.zenml.io \n",
     "\n",
     "import IPython\n",
+    "\n",
     "IPython.Application.instance().kernel.do_shutdown(restart=True)"
    ]
   },
@@ -73,22 +74,16 @@
    "outputs": [],
    "source": [
     "# Do the imports at the top\n",
-    "from typing_extensions import Annotated\n",
-    "from sklearn.datasets import load_breast_cancer\n",
+    "from uuid import UUID\n",
     "\n",
     "import pandas as pd\n",
-    "from zenml import step, pipeline, Model, get_step_context\n",
+    "from pipelines import feature_engineering, training\n",
+    "from sklearn.datasets import load_breast_cancer\n",
+    "from steps import data_loader, inference_preprocessor\n",
+    "from typing_extensions import Annotated\n",
+    "from zenml import Model, get_step_context, pipeline, step\n",
     "from zenml.client import Client\n",
     "from zenml.logger import get_logger\n",
-    "from uuid import UUID\n",
-    "\n",
-    "from zenml import pipeline\n",
-    "\n",
-    "from steps import (\n",
-    "    data_loader,\n",
-    "    inference_preprocessor\n",
-    ")\n",
-    "from pipelines import feature_engineering, training\n",
     "\n",
     "logger = get_logger(__name__)\n",
     "\n",
@@ -126,20 +121,22 @@
     "@step\n",
     "def data_loader_simplified(\n",
     "    random_state: int, is_inference: bool = False, target: str = \"target\"\n",
-    ") -> Annotated[pd.DataFrame, \"dataset\"]:  # We name the dataset \n",
+    ") -> Annotated[pd.DataFrame, \"dataset\"]:  # We name the dataset\n",
     "    \"\"\"Dataset reader step.\"\"\"\n",
     "    dataset = load_breast_cancer(as_frame=True)\n",
     "    inference_size = int(len(dataset.target) * 0.05)\n",
     "    dataset: pd.DataFrame = dataset.frame\n",
-    "    inference_subset = dataset.sample(inference_size, random_state=random_state)\n",
+    "    inference_subset = dataset.sample(\n",
+    "        inference_size, random_state=random_state\n",
+    "    )\n",
     "    if is_inference:\n",
     "        dataset = inference_subset\n",
     "        dataset.drop(columns=target, inplace=True)\n",
     "    else:\n",
     "        dataset.drop(inference_subset.index, inplace=True)\n",
     "    dataset.reset_index(drop=True, inplace=True)\n",
     "    logger.info(f\"Dataset with {len(dataset)} records loaded!\")\n",
-    "    return dataset\n"
+    "    return dataset"
    ]
   },
   {
@@ -243,7 +240,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "feature_engineering(random_state=42,test_size=0.25)"
+    "feature_engineering(random_state=42, test_size=0.25)"
    ]
   },
   {
@@ -340,7 +337,9 @@
    "outputs": [],
    "source": [
     "# Get artifact version from our run\n",
-    "dataset_trn_artifact_version_via_run = run.steps[\"data_preprocessor\"].outputs[\"dataset_trn\"] \n",
+    "dataset_trn_artifact_version_via_run = run.steps[\"data_preprocessor\"].outputs[\n",
+    "    \"dataset_trn\"\n",
+    "]\n",
     "\n",
     "# Get latest version from client directly\n",
     "dataset_trn_artifact_version = client.get_artifact_version(\"dataset_trn\")\n",
@@ -359,7 +358,9 @@
    "source": [
     "# Fetch the rest of the artifacts\n",
     "dataset_tst_artifact_version = client.get_artifact_version(\"dataset_tst\")\n",
-    "preprocessing_pipeline_artifact_version = client.get_artifact_version(\"preprocess_pipeline\")"
+    "preprocessing_pipeline_artifact_version = client.get_artifact_version(\n",
+    "    \"preprocess_pipeline\"\n",
+    ")"
    ]
   },
   {
@@ -480,7 +481,7 @@
     "training.with_options(enable_cache=False)(\n",
     "    model_type=\"xgboost\",\n",
     "    train_dataset_id=dataset_trn_artifact_version.id,\n",
-    "    test_dataset_id=dataset_tst_artifact_version.id\n",
+    "    test_dataset_id=dataset_tst_artifact_version.id,\n",
     ")\n",
     "\n",
     "xgboost_run = client.get_pipeline(\"training\").last_run"
@@ -497,7 +498,7 @@
     "sgd_run = training.with_options(enable_cache=False)(\n",
     "    model_type=\"sgd\",\n",
     "    train_dataset_id=dataset_trn_artifact_version.id,\n",
-    "    test_dataset_id=dataset_tst_artifact_version.id\n",
+    "    test_dataset_id=dataset_tst_artifact_version.id,\n",
     ")\n",
     "\n",
     "sgd_run = client.get_pipeline(\"training\").last_run"
@@ -521,7 +522,9 @@
    "outputs": [],
    "source": [
     "# The evaluator returns a float value with the accuracy\n",
-    "xgboost_run.steps[\"model_evaluator\"].output.load() >= sgd_run.steps[\"model_evaluator\"].output.load()"
+    "xgboost_run.steps[\"model_evaluator\"].output.load() >= sgd_run.steps[\n",
+    "    \"model_evaluator\"\n",
+    "].output.load()"
    ]
   },
   {
@@ -579,7 +582,7 @@
     "training_configured(\n",
     "    model_type=\"sgd\",\n",
     "    train_dataset_id=dataset_trn_artifact_version.id,\n",
-    "    test_dataset_id=dataset_tst_artifact_version.id\n",
+    "    test_dataset_id=dataset_tst_artifact_version.id,\n",
     ")"
    ]
   },
@@ -601,7 +604,7 @@
     "training_configured(\n",
     "    model_type=\"xgboost\",\n",
     "    train_dataset_id=dataset_trn_artifact_version.id,\n",
-    "    test_dataset_id=dataset_tst_artifact_version.id\n",
+    "    test_dataset_id=dataset_tst_artifact_version.id,\n",
     ")"
    ]
   },
@@ -650,10 +653,14 @@
    "outputs": [],
    "source": [
     "# Let's load the XGBoost version\n",
-    "xgboost_zenml_model_version = client.list_model_versions(\"breast_cancer_classifier\", tag=\"xgboost\")[-1]\n",
+    "xgboost_zenml_model_version = client.list_model_versions(\n",
+    "    \"breast_cancer_classifier\", tag=\"xgboost\"\n",
+    ")[-1]\n",
     "\n",
     "# We can now load our classifier directly as well\n",
-    "xgboost_classifier = xgboost_zenml_model_version.get_artifact(\"breast_cancer_classifier\").load()\n",
+    "xgboost_classifier = xgboost_zenml_model_version.get_artifact(\n",
+    "    \"breast_cancer_classifier\"\n",
+    ").load()\n",
     "\n",
     "xgboost_classifier"
    ]
@@ -750,7 +757,9 @@
    "outputs": [],
    "source": [
     "@step\n",
-    "def inference_predict(dataset_inf: pd.DataFrame) -> Annotated[pd.Series, \"predictions\"]:\n",
+    "def inference_predict(\n",
+    "    dataset_inf: pd.DataFrame,\n",
+    ") -> Annotated[pd.Series, \"predictions\"]:\n",
     "    \"\"\"Predictions step\"\"\"\n",
     "    # Get the model\n",
     "    model = get_step_context().model\n",
@@ -761,7 +770,7 @@
     "\n",
     "    predictions = pd.Series(predictions, name=\"predicted\")\n",
     "\n",
-    "    return predictions\n"
+    "    return predictions"
    ]
   },
   {
@@ -788,18 +797,18 @@
     "    random_state = 42\n",
     "    target = \"target\"\n",
     "\n",
-    "    df_inference = data_loader(\n",
-    "        random_state=random_state, is_inference=True\n",
-    "    )\n",
+    "    df_inference = data_loader(random_state=random_state, is_inference=True)\n",
     "    df_inference = inference_preprocessor(\n",
     "        dataset_inf=df_inference,\n",
     "        # We use the preprocess pipeline from the feature engineering pipeline\n",
-    "        preprocess_pipeline=client.get_artifact_version(name_id_or_prefix=preprocess_pipeline_id),\n",
+    "        preprocess_pipeline=client.get_artifact_version(\n",
+    "            name_id_or_prefix=preprocess_pipeline_id\n",
+    "        ),\n",
     "        target=target,\n",
     "    )\n",
     "    inference_predict(\n",
     "        dataset_inf=df_inference,\n",
-    "    )\n"
+    "    )"
    ]
   },
   {
@@ -823,7 +832,7 @@
     "# Lets add some metadata to the model to make it identifiable\n",
     "pipeline_settings[\"model\"] = Model(\n",
     "    name=\"breast_cancer_classifier\",\n",
-    "    version=\"production\", # We can pass in the stage name here!\n",
+    "    version=\"production\",  # We can pass in the stage name here!\n",
     ")"
    ]
   },
@@ -863,7 +872,9 @@
    "outputs": [],
    "source": [
     "# Fetch production model\n",
-    "production_model_version = client.get_model_version(\"breast_cancer_classifier\", \"production\")\n",
+    "production_model_version = client.get_model_version(\n",
+    "    \"breast_cancer_classifier\", \"production\"\n",
+    ")\n",
     "\n",
     "# Get the predictions artifact\n",
     "production_model_version.get_artifact(\"predictions\").load()"
@@ -895,6 +906,7 @@
    "outputs": [],
    "source": [
     "from zenml.client import Client\n",
+    "\n",
     "client = Client()"
    ]
   },
@@ -905,8 +917,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sgd_model_version = client.list_model_versions(\"breast_cancer_classifier\",tag=\"sgd\")[-1]\n",
-    "xgboost_model_version = client.list_model_versions(\"breast_cancer_classifier\",tag=\"xgboost\")[-1]\n",
+    "sgd_model_version = client.list_model_versions(\n",
+    "    \"breast_cancer_classifier\", tag=\"sgd\"\n",
+    ")[-1]\n",
+    "xgboost_model_version = client.list_model_versions(\n",
+    "    \"breast_cancer_classifier\", tag=\"xgboost\"\n",
+    ")[-1]\n",
     "print(f\"SGD version is staged as `{sgd_model_version.stage}`\")\n",
     "print(f\"XGBoost version is staged as `{xgboost_model_version.stage}`\")"
    ]
@@ -975,10 +991,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sgd_clf_metadata = sgd_model_version.get_artifact(\"breast_cancer_classifier\").run_metadata\n",
-    "xgboost_clf_metadata = xgboost_model_version.get_artifact(\"breast_cancer_classifier\").run_metadata\n",
-    "print(f\"SGD{' (production)' if sgd_model_version.stage == 'production' else ''} metrics: train={sgd_clf_metadata['train_accuracy'].value*100:.2f}% test={sgd_clf_metadata['test_accuracy'].value*100:.2f}%\")\n",
-    "print(f\"XGBoost{' (production)' if xgboost_model_version.stage == 'production' else ''} metrics: train={xgboost_clf_metadata['train_accuracy'].value*100:.2f}% test={xgboost_clf_metadata['test_accuracy'].value*100:.2f}%\")"
+    "sgd_clf_metadata = sgd_model_version.get_artifact(\n",
+    "    \"breast_cancer_classifier\"\n",
+    ").run_metadata\n",
+    "xgboost_clf_metadata = xgboost_model_version.get_artifact(\n",
+    "    \"breast_cancer_classifier\"\n",
+    ").run_metadata\n",
+    "print(\n",
+    "    f\"SGD{' (production)' if sgd_model_version.stage == 'production' else ''} metrics: train={sgd_clf_metadata['train_accuracy'].value*100:.2f}% test={sgd_clf_metadata['test_accuracy'].value*100:.2f}%\"\n",
+    ")\n",
+    "print(\n",
+    "    f\"XGBoost{' (production)' if xgboost_model_version.stage == 'production' else ''} metrics: train={xgboost_clf_metadata['train_accuracy'].value*100:.2f}% test={xgboost_clf_metadata['test_accuracy'].value*100:.2f}%\"\n",
+    ")"
    ]
   },
   {
@@ -996,21 +1020,27 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import seaborn as sns\n",
-    "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import seaborn as sns\n",
     "\n",
-    "def plot_confusion_matrix(metadata_pointer, tp: str,ax):\n",
-    "    confusion_matrix = np.array(metadata_pointer[\"confusion_matrix\"].value, dtype=float).reshape((2,2))\n",
+    "\n",
+    "def plot_confusion_matrix(metadata_pointer, tp: str, ax):\n",
+    "    confusion_matrix = np.array(\n",
+    "        metadata_pointer[\"confusion_matrix\"].value, dtype=float\n",
+    "    ).reshape((2, 2))\n",
     "    confusion_matrix /= np.sum(confusion_matrix)\n",
-    "    sns.heatmap(confusion_matrix, annot=True,fmt='.2%',cmap=\"coolwarm\",ax=ax)\n",
+    "    sns.heatmap(\n",
+    "        confusion_matrix, annot=True, fmt=\".2%\", cmap=\"coolwarm\", ax=ax\n",
+    "    )\n",
     "    ax.set_title(f\"{tp} confusion matrix\")\n",
     "    ax.set_ylabel(\"Ground Label\")\n",
     "    ax.set_xlabel(\"Predicted Label\")\n",
     "\n",
-    "fig, ax = plt.subplots(1,2,figsize=(15,4))\n",
-    "plot_confusion_matrix(sgd_clf_metadata, \"SGD\",ax[0])\n",
-    "plot_confusion_matrix(xgboost_clf_metadata, \"RF\",ax[1])"
+    "\n",
+    "fig, ax = plt.subplots(1, 2, figsize=(15, 4))\n",
+    "plot_confusion_matrix(sgd_clf_metadata, \"SGD\", ax[0])\n",
+    "plot_confusion_matrix(xgboost_clf_metadata, \"RF\", ax[1])"
    ]
   },
   {
@@ -1052,7 +1082,7 @@
     "for artifact_name, versions in sgd_model_version.data_artifacts.items():\n",
     "    if versions:\n",
     "        print(f\"Existing version of `{artifact_name}`:\")\n",
-    "        for version_name, artifact_ in  versions.items():\n",
+    "        for version_name, artifact_ in versions.items():\n",
     "            print(version_name, artifact_.data_type.attribute)"
    ]
   },