register_model.py

import mlflow
from mlflow.tracking import MlflowClient
from mlflow.models.signature import infer_signature


# Create a sample test DataFrame
test_data = pd.DataFrame({
    'area': ['Any Region'],
    'user_id': [100]
})

# Initialize MLflow client
client = MlflowClient()
# Create an instance of the RankingModel
model = RankingModel()

# Start an MLflow run
with mlflow.start_run() as run:
    # Call the `predict()` method on the instantiated model with the required arguments
    prediction = model.predict(context=None, model_input=test_data)

    # Infer the signature of the predict function
    signature = infer_signature(test_data, prediction)

    # Log the model artifact to MLflow
    mlflow.pyfunc.log_model(
        artifact_path="sql_model",
        python_model=model,
        input_example=test_data,
        signature=signature
    )

    # Register the model to the model registry
    mv = mlflow.register_model(f'runs:/{run.info.run_id}/sql_model', "sql_model")
    client.transition_model_version_stage(f'sql_model', mv.version, "Production", archive_existing_versions=True)