Module 20: Databricks AI & Machine Learning

mlflow.set_experiment(f"/Users/{spark.sql('SELECT current_user()').first()[0]}/dbx_course_experiment")

with mlflow.start_run(run_name="first_run") as run:
    mlflow.log_param("learning_rate", 0.1)
    mlflow.log_metric("accuracy", 0.91)
    mlflow.log_text("a note about this run", "notes.txt")
    print(run.info.run_id)

from sklearn.datasets import load_diabetes
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

X, y = load_diabetes(return_X_y=True, as_frame=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

mlflow.sklearn.autolog(log_models=False)       # log params and metrics automatically

for depth in (3, 6, 10):
    with mlflow.start_run(run_name=f"rf_depth_{depth}"):
        model = RandomForestRegressor(max_depth=depth, n_estimators=50, random_state=42)
        model.fit(X_train, y_train)
        rmse = mean_squared_error(y_test, model.predict(X_test)) ** 0.5
        mlflow.log_metric("test_rmse", rmse)
        print(depth, round(rmse, 2))

# record the data version used for training
mlflow.set_tag("data_table", "dbx_course.orders")
mlflow.set_tag("data_delta_version", spark.sql("DESCRIBE HISTORY orders_delta").first()["version"])

# a feature table is just a table: customer-level features from the practice data
(spark.table("orders").groupBy("customer_id")
      .agg(F.count("*").alias("orders_total"),
           F.avg("quantity").alias("avg_quantity"),
           F.max("order_ts").alias("last_order_ts"))
      .write.mode("overwrite").saveAsTable("customer_features"))

%sql
ALTER TABLE customer_features ALTER COLUMN customer_id SET NOT NULL;
ALTER TABLE customer_features ADD CONSTRAINT customer_features_pk PRIMARY KEY (customer_id)

from sklearn.metrics import mean_absolute_error, r2_score
import mlflow.sklearn

with mlflow.start_run(run_name="rf_final") as run:
    model = RandomForestRegressor(max_depth=6, n_estimators=100, random_state=42).fit(X_train, y_train)
    pred = model.predict(X_test)
    mlflow.log_metrics({
        "test_rmse": mean_squared_error(y_test, pred) ** 0.5,
        "test_mae":  mean_absolute_error(y_test, pred),
        "test_r2":   r2_score(y_test, pred),
    })
    mlflow.sklearn.log_model(model, artifact_path="model", input_example=X_train.iloc[:2])
    run_id = run.info.run_id
print(run_id)

model_name = f"{catalog}.dbx_course.diabetes_rf"
result = mlflow.register_model(f"runs:/{run_id}/model", model_name)
print(result.version)

from mlflow import MlflowClient
client = MlflowClient()

client.set_registered_model_alias(model_name, "champion", result.version)

loaded = mlflow.pyfunc.load_model(f"models:/{model_name}@champion")
print(loaded.predict(X_test.iloc[:3]))

# calling an endpoint (shape only; needs a real endpoint, host and token)
import requests
resp = requests.post(
    "https://<workspace-url>/serving-endpoints/<endpoint-name>/invocations",
    headers={"Authorization": "Bearer <token>"},
    json={"dataframe_split": {"columns": list(X_test.columns), "data": X_test.iloc[:1].values.tolist()}},
)
print(resp.json())

# the data side: a chunk table (works on Free Edition)
docs = spark.createDataFrame(
    [(1, "Delta tables support time travel with VERSION AS OF."),
     (2, "Auto Loader ingests new files incrementally.")],
    "chunk_id INT, text STRING"
).withColumn("source", F.lit("course_notes"))
docs.write.mode("overwrite").option("delta.enableChangeDataFeed", "true").saveAsTable("doc_chunks")

%sql
-- general purpose: any supported foundation model endpoint (needs access to the endpoint)
SELECT ai_query(
         'databricks-meta-llama-3-3-70b-instruct',
         'Classify the sentiment of this review as positive, negative or neutral: ' || review) AS sentiment,
       review
FROM (VALUES ('Fast delivery, great product'), ('Arrived broken')) AS t(review)

%sql
-- task-specific functions: no model name or prompt needed
SELECT ai_analyze_sentiment('I love this') AS sentiment,
       ai_translate('Good morning', 'nl') AS dutch,
       ai_summarize('Delta Lake adds transactions to data lakes.') AS summary