MLflow and DSPy Tutorial for Beginners

MLFlow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MLflow and DSPy Tutorial for Beginners\n",
    "\n",
    "*Learn to manage ML experiments with MLflow and build modular AI solutions with DSPy in Google Colab.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "Machine Learning projects often involve numerous experiments, parameters, and iterations. **MLflow** is an open-source platform for managing the end-to-end machine learning lifecycle. It helps track experiments, models, and metrics, making it easier to organize and reproduce results. On the other hand, **DSPy** (Declarative Self-improving Python) is an open-source Python framework that allows developers to build language model applications using modular and declarative programming instead of relying on one-off prompting techniques. In simpler terms, DSPy lets you *program* your AI rather than just prompt it, enabling more structured and optimizable AI behaviors.\n",
    "\n",
    "In this tutorial, we will introduce beginners to an AI/ML workflow that combines MLflow and DSPy. You'll learn how to set up your environment, manage sensitive API keys (using Google Colab's Secrets and OpenRouter), organize your code and data, integrate a dataset, train and optimize a model, evaluate its performance, and finally deploy the model. Throughout, we'll highlight best practices for reproducibility, experiment tracking, and modular coding. By the end, you should understand how MLflow can track your machine learning experiments and how DSPy can help build and refine AI model logic as code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User Guide\n",
    "\n",
    "Let's walk through the key steps of our AI/ML workflow using MLflow and DSPy:\n",
    "\n",
    "1. **Installation and Setup** – Install MLflow, DSPy, and other required libraries, and configure our environment.\n",
    "2. **Managing Secrets (API Keys)** – Securely handle API keys (like OpenAI keys) using Google Colab Secrets and optionally OpenRouter for model access.\n",
    "3. **Modular Code Structure** – Write reusable functions and classes to keep code organized and maintainable.\n",
    "4. **File/Folder Organization** – Set up a clear project directory structure to manage datasets, models, and code.\n",
    "5. **Data Integration** – Load the dataset and preprocess it for training.\n",
    "6. **Model Training & Optimization** – Use DSPy to train and optimize AI models, while MLflow tracks the training experiments.\n",
    "7. **Evaluation Metrics** – Assess model performance using key metrics and record these results.\n",
    "8. **Deployment** – Save the trained model and demonstrate how it can be loaded or served for use in a production environment.\n",
    "9. **Best Practices** – Ensure reproducibility, logging, and tracking experiments.\n",
    "\n",
    "Run each code cell in order (this notebook is Colab-ready) and modify the examples to experiment on your own."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "First, we need to install and import the necessary libraries. We will use **MLflow** for experiment tracking and **DSPy** for building and optimizing our AI modules. We'll also use common libraries like **pandas** (for data handling) and **scikit-learn** (for a simple machine learning model in this example)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Install MLflow, DSPy, and other required libraries\n",
    "!pip install -q mlflow dspy pandas scikit-learn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Google Colab Secrets and OpenRouter\n",
    "\n",
    "When working with AI models, you often need API keys (for example, for accessing AI services). It's important **not to hardcode keys** in your notebook. Google Colab provides a **Secrets** feature that allows you to securely store API keys and retrieve them in your code. You can add secrets by clicking on the \"Secrets\" icon in Colab, then entering a name (e.g., `OPENAI_API_KEY`) and its value. Once stored, a secret can be accessed from your notebook.\n",
    "\n",
    "To use a stored secret, Colab offers the `google.colab.userdata` module. For example, you can fetch your API key and set it as an environment variable:\n",
    "```python\n",
    "from google.colab import userdata\n",
    "import os\n",
    "os.environ[\"OPENAI_API_KEY\"] = userdata.get('OPENAI_API_KEY')\n",
    "```\n",
    "\n",
    "Alternatively, you can use Python's `getpass` to input the key at runtime (so it won't be visible in the notebook)."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from google.colab import userdata\n",
    "import os, getpass\n",
    "\n",
    "# Retrieve the API key from Colab Secrets, or prompt for it if not found\n",
    "api_key = userdata.get('OPENAI_API_KEY')\n",
    "if api_key is None or api_key == '':\n",
    "    api_key = getpass.getpass('Enter your API key: ')\n",
    "os.environ['OPENAI_API_KEY'] = api_key\n",
    "\n",
    "print('API key set successfully!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modular Code Structures\n",
    "\n",
    "Writing modular code means breaking down your project into reusable components (functions, classes, modules) instead of writing one long script. This approach makes it easier to read, debug, and reuse code. For instance, you might have one function to preprocess data, another to train a model, and another to evaluate results. This modular structure is especially useful in machine learning projects where you may try multiple approaches. DSPy itself encourages a modular approach by letting you define *modules* for each part of an AI task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## File/Folder Structure\n",
    "\n",
    "Organizing files and directories is important for managing larger projects. Rather than keeping everything in one notebook, you should structure your project repository so that code, data, and results are separated logically. For example, a typical project structure might look like this:\n",
    "\n",
    "```text\n",
    "project_name/\n",
    "├── data/\n",
    "│   ├── train.csv\n",
    "│   └── test.csv\n",
    "├── models/\n",
    "│   └── model.pkl\n",
    "├── notebooks/\n",
    "│   └── exploration.ipynb\n",
    "├── src/\n",
    "│   ├── preprocess.py\n",
    "│   ├── train.py\n",
    "│   └── inference.py\n",
    "└── README.md\n",
    "```\n",
    "\n",
    "In Colab, you can also mount Google Drive or use cloud storage to organize your data and outputs in a similar way."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Integration\n",
    "\n",
    "Now that our environment is ready, let's integrate a dataset for our model. Data integration involves loading the data, exploring it, and preprocessing it so it's suitable for training. Common steps include:\n",
    "- **Loading data** from a file or source (CSV, database, etc.),\n",
    "- **Cleaning** the data (handling missing values, removing duplicates),\n",
    "- **Feature engineering or encoding** (e.g., converting categorical values to numeric),\n",
    "- **Splitting** the data into training and testing sets.\n",
    "\n",
    "For simplicity, we'll use the classic Iris dataset from scikit-learn. This dataset contains flower measurements and species labels for classification. We'll load it, create a pandas DataFrame, and then split it into training and test sets."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Load the Iris dataset and prepare the DataFrame\n",
    "from sklearn.datasets import load_iris\n",
    "import pandas as pd\n",
    "\n",
    "data = load_iris()\n",
    "df = pd.DataFrame(data.data, columns=data.feature_names)\n",
    "df['target'] = data.target\n",
    "print('Dataset sample:')\n",
    "print(df.head())\n",
    "\n",
    "# Split the data into training and testing sets\n",
    "from sklearn.model_selection import train_test_split\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    df[data.feature_names], df['target'], test_size=0.2, random_state=42\n",
    ")\n",
    "print('Train set size:', X_train.shape)\n",
    "print('Test set size:', X_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Training and Optimization\n",
    "\n",
    "With our data ready, the next step is to train a model and track the process. We'll start by training a simple logistic regression model on the Iris dataset using MLflow to log parameters, metrics, and the model artifact. After that, we'll demonstrate how **DSPy** can be used to program and optimize a language-model-based component, showcasing the potential of building and refining custom AI modules."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training a Model with MLflow\n",
    "\n",
    "We'll train a logistic regression model on the Iris training set. MLflow will log important details such as model parameters, evaluation metrics, and the model artifact. This logging makes it easy to reproduce and compare experiments."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import mlflow\n",
    "import sklearn\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "# Start an MLflow experiment run\n",
    "mlflow.set_experiment(\"Iris_Classification\")\n",
    "with mlflow.start_run(run_name=\"Logistic Regression\") as run:\n",
    "    # Initialize and train the model\n",
    "    model = LogisticRegression(max_iter=1000)\n",
    "    model.fit(X_train, y_train)\n",
    "    \n",
    "    # Predict on the test set and calculate accuracy\n",
    "    predictions = model.predict(X_test)\n",
    "    acc = accuracy_score(y_test, predictions)\n",
    "    print('Test Accuracy:', acc)\n",
    "    \n",
    "    # Log parameters and metrics to MLflow\n",
    "    mlflow.log_param('model_type', 'LogisticRegression')\n",
    "    mlflow.log_param('max_iter', 1000)\n",
    "    mlflow.log_param('C', 1.0)\n",
    "    mlflow.log_metric('accuracy', acc)\n",
    "    \n",
    "    # Log the trained model as an artifact\n",
    "    mlflow.sklearn.log_model(model, artifact_path=\"model\")\n",
    "    \n",
    "# Get the run ID for later use (for loading the model)\n",
    "run_id = run.info.run_id\n",
    "print('Run ID:', run_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimizing Model Behavior with DSPy\n",
    "\n",
    "Now we'll showcase DSPy. Instead of focusing solely on traditional machine learning models, DSPy lets us create and optimize AI modules—miniature brains that can understand and execute specific tasks. Rather than relying on prompt engineering alone, DSPy allows you to program AI behavior, enabling deeper customization. In this example, we configure DSPy to use the **o3-mini-high** model (our new choice) and create a simple sentiment analysis module. This model will process an input sentence and output a boolean sentiment indicator.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import dspy\n",
    "\n",
    "# Configure DSPy with the o3-mini-high model (make sure your API key is set)\n",
    "lm = dspy.LM('o3-mini-high', api_key=os.getenv('OPENAI_API_KEY'))\n",
    "dspy.configure(lm=lm)\n",
    "\n",
    "# Define a DSPy prediction module for sentiment analysis\n",
    "classifier = dspy.Predict('sentence -> sentiment: bool')\n",
    "\n",
    "# Test the module with an example sentence\n",
    "sentence = \"I love the new design, it's brilliant!\"\n",
    "result = classifier(sentence=sentence)\n",
    "print('Input:', sentence)\n",
    "print('Predicted sentiment (positive=True):', result.sentiment)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation Metrics\n",
    "\n",
    "Evaluating model performance is a critical part of any ML workflow. For our classifier, we used **accuracy** on the test set as the primary metric (the proportion of correctly predicted samples). Depending on the problem, you might also consider other metrics such as precision, recall, and F1-score for classification tasks, or MSE/MAE for regression tasks. MLflow allows you to log any number of metrics during training and evaluation, making it easier to compare different runs and choose the best model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deployment\n",
    "\n",
    "After training and evaluating a model, the final step is deployment. Deployment can take many forms:\n",
    "- **Batch predictions**: Using the model to generate predictions on new data in batches.\n",
    "- **Real-time serving**: Hosting the model behind an API for on-demand predictions.\n",
    "\n",
    "MLflow Models make deployment easier by packaging the model with its dependencies. In this notebook, we simulate deployment by loading the logged model and running predictions on new data. In a production setting, you might serve the model via MLflow's REST API or deploy it to a cloud platform."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Load the model from the MLflow run and use it for inference\n",
    "import mlflow\n",
    "\n",
    "# Construct the model URI using the run ID and artifact path\n",
    "model_uri = f\"runs:/{run_id}/model\"\n",
    "loaded_model = mlflow.sklearn.load_model(model_uri)\n",
    "\n",
    "# Perform a sample prediction using the loaded model\n",
    "sample_inputs = X_test[:5]\n",
    "sample_true = y_test[:5]\n",
    "predictions = loaded_model.predict(sample_inputs)\n",
    "print('Sample input features:\\n', sample_inputs)\n",
    "print('Model predictions:', predictions.tolist())\n",
    "print('Actual labels:    ', sample_true.tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Best Practices\n",
    "\n",
    "To wrap up, here are some best practices for AI/ML workflows:\n",
    "\n",
    "- **Track Everything**: Use MLflow to log parameters, code versions, data versions, metrics, and artifacts. This makes your work reproducible and easier to debug.\n",
    "- **Reproducibility**: Ensure that your experiments are reproducible by fixing random seeds, pinning package versions, and capturing environment details.\n",
    "- **Version Control**: Keep your code in version control (e.g., Git). Log commit hashes in MLflow for traceability.\n",
    "- **Modular Design**: Write modular code to isolate and reuse components across projects.\n",
    "- **Testing and Validation**: Validate your model and code with appropriate testing methods.\n",
    "- **Model Registry & Deployment**: Manage model versions using a registry and automate deployment when possible.\n",
    "- **Continuous Monitoring**: Once deployed, continuously monitor your model's performance to address data drift or model decay.\n",
    "\n",
    "By following these practices, you create a robust ML workflow. MLflow handles tracking and reproducibility, while DSPy enables you to build and optimize AI modules in a modular, programmable way. Together, these tools empower rapid prototyping and innovation in AI."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}