Salary Prediction with Linear Regression in Python — scikit-learn Train and Predict Example

A beginner-friendly machine learning example in Python: a scikit-learn linear regression pipeline that predicts salary from years of experience, with a separate training script that saves the model and scaler with joblib, and an inference script that loads them to make predictions.

Overview

This project has two scripts:

  • Model Builder/builder.py — cleans data, trains the model, evaluates it, and saves it to disk
  • Prediction/prediction.py — loads the saved model and predicts salary from user input

Scripts

Model Builder (builder.py)

A complete ML pipeline that processes a CSV dataset and produces a trained model.

Pipeline Steps

Step Action
1 Load dataset from Source/Salary.csv
2 Remove duplicate rows
3 Drop columns with >50% missing values; fill remaining nulls with column mean
4 Scale years_experience using StandardScaler
5 Split into 80% train / 20% test sets
6 Train LinearRegression model
7 Evaluate with R² score and Mean Squared Error
8 Save model and scaler if R² > 0.95

Usage

Both scripts open files with paths relative to the working directory (./Source/...), so run them from the Salary folder:

python "Model Builder/builder.py"

Output format:

R² Score: <score to 4 decimal places>
Mean Squared Error: <mse to 2 decimal places>
Model and scaler saved!

If R² ≤ 0.95, the model is not saved and Model not saved as R² score is below 95%. is printed instead.

Salary Predictor (prediction.py)

Loads the saved model and scaler, accepts user input, and outputs a predicted salary.

Usage

From the Salary folder:

python "Prediction/prediction.py"

Session format:

Enter Years of Experience: 7
Predicted Salary for 7.0 years of experience: $<predicted amount>

Requirements

pandas
scikit-learn
joblib
numpy

Install dependencies:

pip install pandas scikit-learn joblib numpy

File Structure

Salary/
├── Model Builder/
│   └── builder.py          # Training pipeline
├── Prediction/
│   └── prediction.py       # Inference script
└── Source/
    ├── Salary.csv           # Training dataset (205 rows: years_experience, salary)
    ├── salary.pkl           # Saved model (included; regenerated by training)
    └── scaler.pkl           # Saved scaler (included; regenerated by training)

Notes

  • A trained salary.pkl and scaler.pkl are already included, so prediction.py works without training first. Run builder.py to retrain them. If your scikit-learn version differs from the one that created the files, retraining avoids version-mismatch warnings when loading them.
  • The dataset must contain years_experience and salary columns.
  • The model is only saved when R² exceeds 0.95 to ensure prediction quality.

Read all Python Projects and Code Examples guides online · Project overview · All developer projects