/sklearnflask

Flask API for training and predicting using scikit learn models

Primary LanguagePythonMIT LicenseMIT

Flask API for scikit learn

A simple Flask application that can serve predictions from a scikit-learn model. Reads a pickled sklearn model into memory when the Flask app is started and returns predictions through the /predict endpoint. You can also use the /train endpoint to train/retrain the model. Any sklearn model can be used for prediction.

Read more in this blog post.

Dependencies

  • scikit-learn
  • Flask
  • pandas
  • numpy
pip install -r requirements.txt

Running API

python main.py <port>

Endpoints

/predict (POST)

Returns an array of predictions given a JSON object representing independent variables. Here's a sample input:

[
    {"Age": 85, "Sex": "male", "Embarked": "S"},
    {"Age": 24, "Sex": "female", "Embarked": "C"},
    {"Age": 3, "Sex": "male", "Embarked": "C"},
    {"Age": 21, "Sex": "male", "Embarked": "S"}
]

and sample output:

{"prediction": [0, 1, 1, 0]}

/train (GET)

Trains the model. This is currently hard-coded to be a random forest model that is run on a subset of columns of the titanic dataset.

/wipe (GET)

Removes the trained model.