Simple tips prediction using linear regression

linear_regression.py

# Import python packages
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
# Load our dataset from a csv
dataset = pd.read_csv('tip.csv')
# Fetch the age column
age = dataset.iloc[:, :-1].values
# Fetch the tip column
tip = dataset.iloc[:, 1].values
# Splitting the data 20% to test data
age_train, age_test, tip_train, tip_test = train_test_split(age, tip, test_size = 0.2, random_state = 0)

# Create an instance of Linear Regression
regressor = LinearRegression()
regressor.fit(age_train, tip_train)

# Let's check how much a 90 years old client will tip
my_test_tip = regressor.predict(np.array(90))

tip_predicted = regressor.predict(age_test)

# We can also visualise the data
# Training set
plt.scatter(age_train, tip_train, color = 'black')
plt.plot(age_train, regressor.predict(age_train), color = 'red')
plt.title('Training set - How much X years old will tip')
plt.xlabel('Age')
plt.ylabel('Tip')
plt.show()
# Test set
plt.scatter(age_test, tip_test, color = 'black')
plt.plot(age_test, regressor.predict(age_test), color = 'red')
plt.title('Test set - How much X years old will tip')
plt.xlabel('Age')
plt.ylabel('Tip')
plt.show()