Tokenize and predict in a single step (deprecates Serving API)

Author: mostafa

api/api.py

@@ -1,41 +1,57 @@
-import os
-
-import pandas as pd
 from flask import Flask, jsonify, request
+import tensorflow as tf
 from tensorflow.keras.preprocessing.sequence import pad_sequences
 from tensorflow.keras.preprocessing.text import Tokenizer
+import pandas as pd
+import os
 
+app = Flask(__name__)
+
+# Constants and configurations
 MAX_WORDS = 10000
 MAX_LEN = 100
 DATASET_PATH = os.getenv("DATASET_PATH", "dataset/sqli_dataset1.csv")
+MODEL_PATH = os.getenv("MODEL_PATH", "/app/sqli_model/3/")
 DATASET = pd.read_csv(DATASET_PATH)
+
+# Tokenizer setup
 TOKENIZER = Tokenizer(num_words=MAX_WORDS, filters="")
 TOKENIZER.fit_on_texts(DATASET["Query"])
-CONFIG = {"DEBUG": False}
 
+# Load the model using tf.saved_model.load and get the serving signature
+loaded_model = tf.saved_model.load(MODEL_PATH)
+model_predict = loaded_model.signatures["serving_default"]
 
-app = Flask(__name__)
-app.config.from_mapping(CONFIG)
 
+@app.route("/predict", methods=["POST"])
+def predict():
+    if not request.json or "query" not in request.json:
+        return jsonify({"error": "No query provided"}), 400
+
+    try:
+        # Tokenize and pad the input query
+        query = request.json["query"]
+        query_seq = TOKENIZER.texts_to_sequences([query])
+        query_vec = pad_sequences(query_seq, maxlen=MAX_LEN)
 
-@app.route("/tokenize_and_sequence", methods=["POST"])
-def tokenize_and_sequence():
-    """Tokenize and sequence the input query from the request
-    and return the vectorized output.
-    """
+        # Convert input to tensor
+        input_tensor = tf.convert_to_tensor(query_vec, dtype=tf.float32)
 
-    body = request.get_json()
-    if not body:
-        return jsonify({"error": "No JSON body provided"}), 400
+        # Use the loaded model's serving signature to make the prediction
+        prediction = model_predict(input_tensor)
 
-    # Vectorize the sample
-    query_seq = TOKENIZER.texts_to_sequences([body["query"]])
-    query_vec = pad_sequences(query_seq, maxlen=MAX_LEN)
+        if "output_0" not in prediction or prediction["output_0"].get_shape() != [1, 1]:
+            return jsonify({"error": "Invalid model output"}), 500
 
-    tokens = query_vec.tolist()
-    return jsonify({"tokens": tokens[0]})
+        return jsonify(
+            {
+                "confidence": float("%.4f" % prediction["output_0"].numpy()[0][0]),
+            }
+        )
+    except Exception as e:
+        # TODO: Log the error and return a proper error message
+        return jsonify({"error": str(e)}), 500
 
 
 if __name__ == "__main__":
-    # Run the app in debug mode
-    app.run(host="localhost", port=8000, debug=True)
+    app.run(host="0.0.0.0", port=8000, debug=True)