chore: import upstream snapshot with attribution

examples/statsmodels/train.py

@@ -0,0 +1,55 @@
+import argparse
+
+import numpy as np
+import statsmodels.api as sm
+from sklearn.metrics import mean_squared_error
+
+import mlflow
+import mlflow.statsmodels
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Statsmodels example")
+    parser.add_argument(
+        "--inverse-method",
+        type=str,
+        default="pinv",
+        help="Can be 'pinv', or 'qr'. 'pinv' uses the Moore-Penrose pseudoinverse "
+        "to solve the least squares problem. 'qr' uses the QR factorization. "
+        "(default: 'pinv')",
+    )
+    return parser.parse_args()
+
+
+def main():
+    # parse command-line arguments
+    args = parse_args()
+
+    # prepare train and test data
+    # Ordinary Least Squares (OLS)
+    np.random.seed(9876789)
+    nsamples = 100
+    x = np.linspace(0, 10, 100)
+    X = np.column_stack((x, x**2))
+    beta = np.array([1, 0.1, 10])
+    e = np.random.normal(size=nsamples)
+    X = sm.add_constant(X)
+    y = np.dot(X, beta) + e
+
+    # enable auto logging
+    mlflow.statsmodels.autolog()
+
+    with mlflow.start_run():
+        ols = sm.OLS(y, X)
+        model = ols.fit(method=args.inverse_method)
+
+        # evaluate model
+        y_pred = model.predict(X)
+        mse = mean_squared_error(y, y_pred)
+
+        # log metrics
+        mlflow.log_metrics({"mse": mse})
+
+
+if __name__ == "__main__":
+    main()