Update app.py

app.py

@@ -1,16 +1,6 @@
-import os
-import requests
-import json
-from io import BytesIO
-
-from flask import Flask, jsonify, render_template, request, send_file
-
-from modules.inference import infer_t5
-from modules.dataset import query_emotion
-
-# https://huggingface.co/settings/tokens
-# https://huggingface.co/spaces/{username}/{space}/settings
-API_TOKEN = os.getenv("BIG_GAN_TOKEN")
+import os, pickle
+import logging
+from flask import Flask, jsonify, render_template, request
 
 app = Flask(__name__)
 
@@ -19,42 +9,43 @@ app = Flask(__name__)
 def index():
     return render_template("index.html")
 
-
-@app.route("/infer_biggan")
-def biggan():
-    input = request.args.get("input")
-
-    output = requests.request(
-        "POST",
-        "https://api-inference.huggingface.co/models/osanseviero/BigGAN-deep-128",
-        headers={"Authorization": f"Bearer {API_TOKEN}"},
-        data=json.dumps(input),
-    )
-
-    return send_file(BytesIO(output.content), mimetype="image/png")
-
-
-@app.route("/infer_t5")
-def t5():
-    input = request.args.get("input")
-
-    output = infer_t5(input)
-
-    return jsonify({"output": output})
-
-
-@app.route("/query_emotion")
-def emotion():
-    start = request.args.get("start")
-    end = request.args.get("end")
-
-    print(start)
-    print(end)
-
-    output = query_emotion(int(start), int(end))
-
-    return jsonify({"output": output})
+@app.route('/api/predict', methods=['POST'])
+def predict():
+  try:
+    # load and parse input
+    data = request.json
+    vector = [
+        float(data['age']),
+        data['travel'],
+        data['department'],
+        float(data['distance']),
+        float(data['education']),
+        data['gender'],
+        float(data['satisfaction']),
+        data['maritalstatus'],
+        float(data['income']),
+        data['overtime'],
+        float(data['totalyears']),
+        float(data['years']),
+        float(data['lastpromotion'])
+    ]
+    # app.logging.info(f'vector: {vector}')
+    # print(f'vector: {vector}\n')
+
+    # load the model
+    with open(os.path.join('data', 'logistic.pkcls'), 'rb') as file:
+      model = pickle.load(file)
+
+    predictions = model(vector, 1)
+    # app.logging.info(f'predictions: {predictions}')
+    # print(f'predictions: {predictions}\n')
+
+    # send the response
+    return jsonify({ "predictions": { "leave": predictions[0], "stay": predictions[1] } })
+  except Exception as e:
+    return jsonify({"error": repr(e)})
 
 
 if __name__ == "__main__":
+    app.logger.setLevel(logging.INFO)
     app.run(host="0.0.0.0", port=7860)