Skip to content
This repository was archived by the owner on Mar 25, 2026. It is now read-only.

Latest commit

 

History

History
23 lines (17 loc) · 1.1 KB

File metadata and controls

23 lines (17 loc) · 1.1 KB

Step 5: Setting up the Databricks configuration

The total distance and baseline premium is sent to Databricks for ML prediction. For ease of purposes, we’ve included the API call that easily allows you to connect to your ML model. Although we’ve decided to use Databricks, you are free to use any ML platform of your choice.

Here's the code we used for the Databricks model (extremely basic), the following code also includes the mlflow commands to post it as an experiment so it can be subsequently deployed as an endpoint:

import mlflow

class MyModelWBaseline(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        return self.my_custom_function(model_input)

    def my_custom_function(self, model_input):
        # do something with the model input
        return model_input[0] + model_input[1]*0.1
        
# save the model
my_model = MyModelWBaseline()
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(artifact_path="model", python_model=my_model)

In Step 6, we will write the ML prediction to MongoDB.