upd/add: pred_proba & refactored

api_server.py

@@ -1,6 +1,12 @@
-"""
-official fastapi HF example https://huggingface.co/docs/hub/spaces-sdks-docker-examples#docker-spaces-examples
-"""
+# official fastapi HF example https://huggingface.co/docs/hub/spaces-sdks-docker-examples#docker-spaces-examples
+
+##################
+# Flask API usages:
+# 1. Just a wrapper over OpenAI API
+# 2. You can use Chain calls of OpenAI API
+# 3. Using your own ML model in combination with openAPI functionality
+# 4. ...
+##################
 
 import os
 import time
@@ -18,26 +24,27 @@ from flask import Flask, jsonify, request, render_template
 load_type = 'remote_hub_from_pretrained'
 """
 local;
-remote_hub_download; - /cache error even using TRANSFORMERS_CACHE & cache_dir to local folder
-remote_hub_from_pretrained; - /cache error even using TRANSFORMERS_CACHE & cache_dir to local folder
+remote_hub_download; 
+remote_hub_from_pretrained;
 remote_hub_pipeline; - needs config.json and this is not easy to grasp how to do it with custom models
 https://discuss.huggingface.co/t/how-to-create-a-config-json-after-saving-a-model/10459/4
 """
 
 REPO_ID = "1vash/mnist_demo_model"
+MODEL_DIR = "./artifacts/models"
 
 # Load the saved model into memory
 if load_type == 'local':
-    model = keras.models.load_model('artifacts/models/mnist_model.h5')
+    model = keras.models.load_model(f'{MODEL_DIR}/mnist_model.h5')
 elif load_type == 'remote_hub_download':
     from huggingface_hub import hf_hub_download
 
     model = keras.models.load_model(hf_hub_download(repo_id=REPO_ID, filename="saved_model.pb"))
 elif load_type == 'remote_hub_from_pretrained':
     # https://huggingface.co/docs/hub/keras
-    os.environ['TRANSFORMERS_CACHE'] = str(Path('./artifacts/').absolute())
+    os.environ['TRANSFORMERS_CACHE'] = str(Path(MODEL_DIR).absolute())
     from huggingface_hub import from_pretrained_keras
-    model = from_pretrained_keras(REPO_ID, cache_dir='./artifacts/')
+    model = from_pretrained_keras(REPO_ID, cache_dir=MODEL_DIR)
 elif load_type == 'remote_hub_pipeline':
     from transformers import pipeline
 
@@ -62,7 +69,8 @@ def predict():
 
     Response format:
     {
-        "prediction": predicted_label,
+        "label": predicted_label,
+        "pred_proba" prediction class probability
         "ml-latency-ms": latency_in_milliseconds
             (Measures time only for ML operations preprocessing with predict)
     }
@@ -78,9 +86,6 @@ def predict():
     # Get pixels out of file
     image_data = Image.open(file)
 
-    # Get the image data from the request
-    # image_data = request.get_json()['image']
-
     # Preprocess the image
     processed_image = preprocess_image(image_data)
 
@@ -89,13 +94,17 @@ def predict():
 
     # Get the predicted class label
     predicted_label = np.argmax(prediction)
+    proba = prediction[0][predicted_label]
 
     # Calculate latency in milliseconds
     latency_ms = (time.time() - start_time) * 1000
 
-    # Return the prediction result and latency as JSON response
-    response = {'prediction': int(predicted_label),
-                'ml-latency-ms': round(latency_ms, 4)}
+    # Return the prediction result and latency as dictionary response
+    response = {
+        'label': int(predicted_label),
+        'pred_proba': float(proba),
+        'ml-latency-ms': round(latency_ms, 4)
+    }
 
     # dictionary is not a JSON: https://www.quora.com/What-is-the-difference-between-JSON-and-a-dictionary
     # flask.jsonify vs json.dumps https://sentry.io/answers/difference-between-json-dumps-and-flask-jsonify/
@@ -149,11 +158,3 @@ def hello_world():
 # Start the Flask application
 if __name__ == '__main__':
     app.run(debug=True)
-
-##################
-# Flask API usages:
-# 1. Just a wrapper over OpenAI API
-# 2. You can use Chain calls of OpenAI API
-# 3. Using your own ML model in combination with openAPI functionality
-# 4. ...
-##################