File size: 2,120 Bytes
4de73ff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import mlflow
import mlflow.sklearn
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
from pyngrok import ngrok
import gradio as gr
# MLflow setup
mlflow.set_tracking_uri("./mlruns") # Local directory for tracking
mlflow.set_experiment("House Price Prediction")
# Training function
def train_and_log_model():
# Load dataset
data = load_diabetes()
X = data.data
y = data.target
# Split dataset
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# Train model
model = LinearRegression()
model.fit(X_train, y_train)
# Predict and evaluate
y_pred = model.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
# Log to MLflow
with mlflow.start_run():
mlflow.log_param("model", "Linear Regression")
mlflow.log_metric("mse", mse)
mlflow.sklearn.log_model(model, "model")
return mse, "Model training complete and logged to MLflow!"
# Start MLflow UI with Ngrok
def start_mlflow_ui():
public_url = ngrok.connect(5000) # Expose the MLflow UI running on port 5000
mlflow_command = "mlflow ui --host 0.0.0.0 --port 5000"
return_code = os.system(mlflow_command)
if return_code != 0:
return "Error: Unable to start MLflow UI."
return f"MLflow UI is accessible at {public_url}"
# Gradio Interface Functions
def train_model():
mse, message = train_and_log_model()
return f"MSE: {mse}\n{message}"
def get_mlflow_ui_link():
public_url = start_mlflow_ui()
return public_url
# Gradio UI
with gr.Blocks() as demo:
gr.Markdown("## House Price Prediction with MLflow")
train_btn = gr.Button("Train Model and Log to MLflow")
mlflow_btn = gr.Button("Start MLflow UI")
output = gr.Textbox(label="Output")
train_btn.click(train_model, inputs=[], outputs=output)
mlflow_btn.click(get_mlflow_ui_link, inputs=[], outputs=output)
# Launch Gradio App
if __name__ == "__main__":
demo.launch()
|