How to Build a REST API for Your Machine Learning Model Using Flask

Muhammad Bilal Shinwari
Red Buffer
Published in
8 min readMar 20, 2024

Machine learning models are now essential components of today’s applications, used in things like suggesting content or recognizing images. To allow other software systems to use these models, you can create a REST API. In this article, we will look at how to make a RESTful API for a machine learning model using Flask, a popular web framework for Python.

In this article, we’ll use a straightforward linear regression algorithm with scikit-learn to keep things simple.

Prerequisites

  1. Python: Make sure you have Python installed on your system. You can download it from the official website: Python Downloads.
  2. Flask: We’ll be using Flask, a lightweight Python web framework, to create our API. You can install it using pip:
pip install flask

3. scikit-learn: For our machine learning model, we’ll use scikit-learn, a powerful library for machine learning in Python. You can install it using pip as well:

pip install scikit-learn

4. Pandas Library: We will use this library for loading our csv files and working with DataFrames . You can install it using pip :

pip install pandas

5. Requests Library: To test our Flask API using the requests library in Python. You can install it using pip as well:

pip install requests

We will need to create following files

  1. model.py
  2. model.pkl
  3. app.py
  4. request.py

model.py: This file is where we create and train your machine learning model. After training, we use the pickle library to save the model to a file named model.pkl. Think of it like saving our model on our computer so we can use it later.

app.py: In this file, we build a REST API using Flask. This API will handle incoming requests, particularly POST requests. When these requests come in, our API will load the trained model from model.pkl using pickle.load(). It's like taking our saved model and putting it to work.

request.py: This file is responsible for sending requests to our API. we send information (called “features”) to our API, and it uses the loaded model to make predictions. We will then receive the results from our API.

Coding Part

  1. model.py

In this file we will develop and train our machine learning model and train it. and we will predict Yearly Amount Spent for E-Commerce Customers based on their time spent on APP, Website etc. You can access the dataset here.

First, we will import libraries. from developing and taring our model.

import pandas as pd
import pickle
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

Importing these libraries will help us work with data and build machine learning models in Python. Pandas organizes data like a virtual spreadsheet, Pickle saves and loads Python things (like models), and Scikit-Learn gives you tools for making predictions from data, such as creating a linear regression model.

df=pd.read_csv('Ecommerce Customers')
features=['Avg. Session Length', 'Time on App','Time on Website',
'Length of Membership']
label="Yearly Amount Spent"
X=df[features]
y=df[label]

we’ve loaded a dataset (‘Ecommerce Customers’) into a Pandas DataFrame called df .We've chosen specific features ('Avg. Session Length', 'Time on App', 'Time on Website', 'Length of Membership') that we want to use for predictions, and we've designated 'Yearly Amount Spent' as the target we want to predict. By placing these features in X and the target in y, we've prepared our data for machine learning, with X holding input features and y containing the output labels for training a predictive model.

X_train,X_test,y_train,y_test=train_test_split(X, y, test_size=0.4,
random_state=42)

This line of code is splitting our data into two parts: X_train and X_test for the input features and y_train and y_test for the output labels. We’re using 40% of the data for testing and 60% for training our machine learning model. The random_state=42 ensures that the split is consistent each time we run the code, making our results reproducible. This splitting is crucial for evaluating and training our model effectively.

Regression_model = LinearRegression()
Regression_model.fit(X_train,y_train)

predication=Regression_model.predict(X_test)

In these lines, we’re creating a linear regression model and training it with the X_train and y_train data. Then, we’re using the trained model to make predictions on the X_test data, and the results are stored in the predication variable. This allows us to see how well our model can predict output values on new, unseen data.

pickle.dump(Regression_model,open("model.pkl","wb"))

In this line of code, we’re using the pickle library to save our trained Regression_model (the linear regression model) to a file named model.pkl in binary write mode (“wb”). This file will contain a serialized version of the trained model, which can be loaded and reused later for making predictions without having to retrain the model from scratch.

And we will load this model in our App.py file in our REST API to make prediction of the requests using pickle.load() function.

So, our model.py is ready to train and save the model and down below is the full code of it.

# Import necessary libraries
import pandas as pd # For data handling
import pickle # For saving the trained model
from sklearn.model_selection import train_test_split # For splitting data
from sklearn.linear_model import LinearRegression #For Fitting the model

# Load the dataset from a CSV file
df = pd.read_csv('Ecommerce Customers.csv')

# Define the features (input) and label (output) columns
features = ['Avg. Session Length', 'Time on App', 'Time on Website',
'Length of Membership']
label = "Yearly Amount Spent"

# Extract input features (X) and output labels (y)
X = df[features]
y = df[label]

# Split the data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.4,
random_state=42)

# Create a Linear Regression model
Regression_model = LinearRegression()

# Train the model on the training data
Regression_model.fit(X_train, y_train)

# Make predictions using the trained model
predictions = Regression_model.predict(X_test)

# Print the model's predictions
print(predictions)

# Save the trained model to a file named "model.pkl"
pickle.dump(Regression_model, open("model.pkl", "wb"))

2. app.py

In this file now we will make our REST API using flask to handle the POST requests for the requests.py or Postman.

import pandas as pd
import pickle
from flask import Flask,render_template,request,jsonify

We are importing pandas (as `pd`) for data manipulation and analysis, pickle for object deserialization, and Flask along with its components like request, and jsonify for web development, enabling the creation of web applications and APIs.

from flask import Flask, request, jsonify

app = Flask(__name__)

# Load your machine learning model here

if __name__ == '__main__':
app.run(debug=True)

This code creates a Flask web application where we can load a machine learning model. When we run the app, it will be in debug mode, helping us identify and fix issues during development. This forms the basis for building a web interface for our machine learning model.

model=pickle.load(open("model.pkl","rb"))

The line of code we provided loads a machine learning model from a file named model.pkl using the Python pickle library. The `”rb”` argument indicates that the file is being opened in binary read mode. Once loaded, the model is stored in the model variable, then we can use it for making predictions. This is a common way to persist and later retrieve trained machine learning models for reuse or deployment.

@app.route("/predict",methods=["POST"])
def predict():
json_=request.json
df=pd.DataFrame(json_)
prediction=model.predict(df)
return jsonify({"Prediction":list(prediction)})

In this code we define route `/predict` in a Flask web application that accepts HTTP POST requests. When a POST request is received at this route, it expects a JSON payload containing data. It then converts this JSON data into a DataFrame using pandas (`df=pd.DataFrame(json_)`).

Next, it uses a machine learning model (presumably loaded previously) to make predictions on the data stored in the DataFrame (`prediction=model.predict(df)`). These predictions are then returned as a JSON response using Flask’s `jsonify` function.

The response contains a key “Prediction” with a list of prediction values. This code essentially creates an endpoint for making predictions using the machine learning model via HTTP POST requests.

So, our Flask API is ready to serve the requests. down below is the full code of app.py file.

import pandas as pd
import pickle
from flask import Flask, request, jsonify

# Creating a Flask app
app = Flask(__name__)

# Loading the machine learning model from a pickle file
model = pickle.load(open("model.pkl", "rb"))

# Define a route for making predictions
@app.route("/predict", methods=["POST"])
def predict():
# Get JSON data from the request
json_ = request.json

# Convert JSON data into a DataFrame
df = pd.DataFrame(json_)

# Use the loaded model to make predictions on the DataFrame
prediction = model.predict(df)

# Return the predictions as a JSON response
return jsonify({"Prediction": list(prediction)})

# Run the Flask app when this script is executed
if __name__ == "__main__":
app.run(debug=True)

Testing Of REST API

We can test our API with request.py as I mentioned that we will send request to the server app.py to make predictions for us.

Here is the full code of request.py which is going to request the sever for predictions.

import requests

# Define the URL of your Flask API
url = 'http://127.0.0.1:5000/predict'

# Define the input data as a dictionary
data = {
"Avg. Session Length": [34.49726773, 31.92627203, 33.00091476, 34.30555663],
"Time on App": [12.65565115, 11.10946073, 11.33027806, 13.71751367],
"Time on Website": [50.57766802, 80.26895887, 37.11059744, 36.72128268],
"Length of Membership": [1.082620633, 2.664034182, 4.104543202, 3.120178783]
}

# Send a POST request to the API with the input data
response = requests.post(url, json=data)

# Check the HTTP response status code
if response.status_code == 200:
# Parse and print the JSON response (assuming it contains the prediction)
prediction = response.json()
print(prediction)
else:
# Handle the case where the API request failed
print(f'API Request Failed with Status Code: {response.status_code}')
print(f'Response Content: {response.text}')

We use the requests library to send a POST request to our Flask-based REST API, accessible at the specified URL. We’ve provided sample input data in JSON format. If the API responds with a 200 status code, we parse and display the JSON response, assuming it contains a prediction. In the event of a failure, we print the status code and response content, enabling us to test the API’s behavior with the given data.

We can also test our API with Postman.

Postman is a widely-used API testing and development tool that simplifies the process of testing and interacting with RESTful APIs. To test a REST API using Postman, follow these steps:

  1. Install Postman: Download and install Postman from the official website.
  2. Open Postman: Launch Postman after installation.
  3. Create a New Request: Click on “New” to create a new request. Give it a name and select the HTTP method (e.g., GET or POST).
  4. Specify the URL: In the request, provide the URL of the API endpoint you want to test.
  5. Set Request Parameters: Depending on your API’s requirements, configure headers, authentication, and request body. For POST requests, use the “Body” tab to define input data.
  6. Send the Request: Click the “Send” button to send the request to the API.
  7. Inspect the Response: Postman will display the response, including the status code, headers, and response body, allowing you to test and verify the API’s functionality.

Conclusion

In summary, our project includes three important files: model.py to train and save our machine learning model, app.py to handle requests, and request.py to make requests to the server. To make everything work smoothly, follow this order: first, run model.py to train the model, then, in a different terminal or window, start app.py to handle requests, and finally, run request.py to send requests and get results from the server. This step-by-step process ensures that the model is ready, the server is running, and requests are processed correctly, making your machine learning model accessible through the RESTful API.

Find all the code in GitHub Repository, FLASK-REST-API-ML

--

--