add: hugging_face model download

api_server.py

@@ -1,6 +1,7 @@
 import os
 import time
 import numpy as np
+from huggingface_hub import hf_hub_download
 
 # Disable tensorflow warnings
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
@@ -8,8 +9,16 @@ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 from tensorflow import keras
 from flask import Flask, jsonify, request
 
+LOCAL = False
+
 # Load the saved model into memory
-model = keras.models.load_model('artifacts/models/mnist_model.h5')
+if LOCAL is True:
+    REPO_ID = "1vash/mnist-model"
+    FILENAME = "mnist_model.h5"
+    model = keras.models.load_model(hf_hub_download(repo_id=REPO_ID, filename=FILENAME))
+else:
+    model = keras.models.load_model('artifacts/models/mnist_model.h5')
+
 
 # Initialize the Flask application
 app = Flask(__name__)

requirements.txt

@@ -3,3 +3,4 @@ flask
 numpy
 requests
 matplotlib
+huggingface_hub