train_linear_regression_model
Train a linear regression model by specifying the output column. Automates data preprocessing, model training, and evaluation, returning the RMSE value for model accuracy.
Instructions
This function trains linear regression model.
Args: Takes input for output column name.
Returns: String which contains the RMSE value.
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| output_column | Yes |
Implementation Reference
- server.py:122-164 (handler)The handler function for the 'train_linear_regression_model' tool, decorated with @mcp.tool(). It trains a linear regression model on data stored in a global context, using the specified output_column as the target, performs train-test split (90/10), fits the model, computes RMSE on test set, and returns the result.@mcp.tool() def train_linear_regression_model(output_column: str) -> str: """ This function trains linear regression model. Args: Takes input for output column name. Returns: String which contains the RMSE value. """ try: data = context.get_data() # Check if the output column exists in the dataset if output_column not in data.columns: return f"Error: '{output_column}' column not found in the dataset." # Prepare the features (X) and target variable (y) X = data.drop(columns=[output_column]) # Drop the target column for features y = data[output_column] # The target variable is the output column # Split the data into training and test sets (80% train, 20% test) X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=42) # Initialize the Linear Regression model model = LinearRegression() # Train the model model.fit(X_train, y_train) # Predict on the test set y_pred = model.predict(X_test) # Calculate RMSE (Root Mean Squared Error) rmse = np.sqrt(mean_squared_error(y_test, y_pred)) # Return the RMSE value return f"Model trained successfully. RMSE: {rmse:.4f}" except Exception as e: return f"An error occurred while training the model: {str(e)}"