Merge pull request #107 from oslabs-beta/yiqun/new

feat: add lib, revise gitignore, revise example

Author: austinbfraser

mlflow/.gitignore

@@ -1,6 +1,12 @@
 # Dependencies
 node_modules/
 
+# Environment variables
+.env
+
+# IDE
+.vscode
+
 # Misc
 .DS_Store
 ../.DS_Store
@@ -9,8 +15,12 @@ node_modules/
 .venv/
 mlruns/
 
-# Build output
-lib/
-
 # Temporary files
-temp/
+temp/
+
+# Test coverage 
+coverage/
+
+
+
+=

mlflow/examples/LinearRegressionExample.js

@@ -1,7 +1,7 @@
 /**
  * @fileoverview Example of using MLflow.js for machine learning projects
- * The example creates a synthetic dataset and trains a linear regression model using TensorFlow.js,
- * tracking various metrics and parameters throughout the training process with MLflow tracking server.
+ * The example creates a synthetic dataset and trains a linear regression model using TensorFlow.js, 
+ * tracking various metrics, hyperparameters and other meta data throughout the training process with MLflow tracking server.
  * It showcases:
  * - Experiment and run management
  * - Hyperparameter logging
@@ -12,19 +12,28 @@
  * - Artifact storage
  *
  * @requires @tensorflow/tfjs-node
+ * @requires @mlflow.js
  *
- * @note Ensure MLflow server is running at http://localhost:5001 before executing
+ * @note Ensure MLflow server is running before executing
  */
 
 import * as tf from '@tensorflow/tfjs-node';
-import MLflow from 'mlflow-js';
+import Mlflow from 'mlflow-js';
 import { fileURLToPath } from 'url';
-import { dirname } from 'path';
+import { dirname, join } from 'path';
+import dotenv from 'dotenv';
 
-const mlflow = new MLflow('http://localhost:5001');
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = dirname(__filename);
+dotenv.config({ path: join(__dirname, '../.env') });
+
+const MLFLOW_TRACKING_URI =
+  process.env.MLFLOW_TRACKING_URI || 'http://localhost:5001';
+const mlflow = new Mlflow(MLFLOW_TRACKING_URI);
+console.log('MLflow server URL:', MLFLOW_TRACKING_URI);
 
 const HYPERPARAMETERS = {
-  learningRate: 0.1,
+  learningRate: 0.25,
   epochs: 10,
   batchSize: 32,
   validationSplit: 0.2,
@@ -168,13 +177,22 @@ async function main() {
     console.log(`MLflow Experiment ID: ${experimentId}`);
 
     // Create run
-    const run = await mlflow.createRun(
-      experimentId,
-      'Simple Linear Regression'
-    );
+    const run = await mlflow.createRun(experimentId);
     const runId = run.info.run_id;
     console.log(`MLflow Run ID: ${runId}`);
 
+    // Generate and prepare data
+    console.log('Generating data...');
+    const { xTrain, yTrain, xTest, yTest } = generateData();
+
+    // Log dataset info
+    const datasetInfo = {
+      dataset: {
+        name: 'synthetic_linear_regression_data',
+      },
+    };
+    await mlflow.logInputs(runId, [datasetInfo]);
+
     // Log hyperparameters
     const params = [
       { key: 'learning_rate', value: HYPERPARAMETERS.learningRate.toString() },
@@ -189,18 +207,6 @@ async function main() {
     ];
     await mlflow.logBatch(runId, undefined, params);
 
-    // Log dataset info
-    const datasetInfo = {
-      dataset: {
-        name: 'synthetic_linear_regression_data',
-      },
-    };
-    await mlflow.logInputs(runId, [datasetInfo]);
-
-    // Generate and prepare data
-    console.log('Generating data...');
-    const { xTrain, yTrain, xTest, yTest } = generateData();
-
     // Train model
     console.log('Creating and training model...');
     const model = createModel();
@@ -209,12 +215,6 @@ async function main() {
     // Evaluate model
     const evaluation = evaluateModel(model, xTest, yTest);
 
-    // Save model artifects
-    const __filename = fileURLToPath(import.meta.url);
-    const __dirname = dirname(__filename);
-    const artifactsPath = `${__dirname}/../mlruns/${experimentId}/${runId}/artifacts`;
-    await model.save(`file://${artifactsPath}`);
-
     // Log evaluation metrics
     const finalMetrics = [
       { key: 'test_mse', value: evaluation.metrics.mse, timestamp: Date.now() },
@@ -235,6 +235,12 @@ async function main() {
     ];
     await mlflow.logBatch(runId, undefined, undefined, paramTags);
 
+    // Save model artifects
+    const __filename = fileURLToPath(import.meta.url);
+    const __dirname = dirname(__filename);
+    const artifactsPath = `${__dirname}/../mlruns/${experimentId}/${runId}/artifacts`;
+    await model.save(`file://${artifactsPath}`);
+
     // Register model if performance meets threshold
     if (evaluation.metrics.r2 > 0.9) {
       const modelName = 'LinearRegression';
@@ -259,7 +265,10 @@ async function main() {
           `runs:/${runId}/model`,
           runId,
           [
-            { key: 'r2', value: evaluation.metrics.r2.toString() },
+            {
+              key: 'learning_rate',
+              value: HYPERPARAMETERS.learningRate.toString(),
+            },
             { key: 'rmse', value: evaluation.metrics.rmse.toString() },
           ]
         );
@@ -277,13 +286,6 @@ async function main() {
       }
     }
 
-    // Log additional metadata
-    const tags = [
-      { key: 'model_type', value: 'linear_regression' },
-      { key: 'data_source', value: 'synthetic' },
-    ];
-    await mlflow.logBatch(runId, undefined, undefined, tags);
-
     // Finish run
     await mlflow.updateRun(runId, 'FINISHED');
 

mlflow/examples/NeuralNetworkHyperparameterTuning.js

@@ -7,9 +7,9 @@ const mlflow = new MLflow('http://localhost:5001');
 
 const HYPERPARAMETER_SPACE = {
   networkArchitectures: [
-    [16, 8], // Small network
-    [32, 16], // Medium network
-    [64, 32], // Larger network
+    [16, 8],
+    [32, 16],
+    [64, 32],
   ],
   learningRates: [0.001, 0.01],
   batchSizes: [32, 64],
@@ -23,7 +23,7 @@ const TRAINING_CONFIG = {
   datasetSize: 2000,
   inputFeatures: 5,
   outputClasses: 3,
-  minibatchSize: 128, // Added for faster training
+  minibatchSize: 128,
 };
 
 // Data generation
@@ -252,7 +252,7 @@ async function main() {
   try {
     console.time('Total Execution Time');
 
-    const experimentName = 'Neural_Network_Hyperparameter_Tuning_Fast';
+    const experimentName = 'Neural_Network_Hyperparameter_Tuning';
     let experimentId;
     try {
       const experiment = await mlflow.getExperimentByName(experimentName);

mlflow/lib/index.d.ts

@@ -0,0 +1,2 @@
+import Mlflow from './mlflow.js';
+export default Mlflow;

mlflow/lib/index.js

@@ -0,0 +1,20 @@
+import ExperimentClient from './tracking/ExperimentClient.js';
+import ExperimentManager from './workflows/ExperimentManager.js';
+import RunClient from './tracking/RunClient.js';
+import RunManager from './workflows/RunManager.js';
+import ModelRegistryClient from './model-registry/ModelRegistryClient.js';
+import ModelVersionClient from './model-registry/ModelVersionClient.js';
+import ModelManager from './workflows/ModelManager.js';
+declare class Mlflow {
+    private components;
+    constructor(trackingUri: string);
+    private initializeMethods;
+    getExperimentClient(): ExperimentClient;
+    getRunClient(): RunClient;
+    getModelRegistryClient(): ModelRegistryClient;
+    getModelVersionClient(): ModelVersionClient;
+    getExperimentManager(): ExperimentManager;
+    getRunManager(): RunManager;
+    getModelManager(): ModelManager;
+}
+export default Mlflow;