import os
import json
from flask import Flask, request, make_response
from logzero import logger

# current_path = os.path.dirname(os.path.abspath(__file__)) # for local
current_path = os.getcwd()  # for docker
logger.info(f"{current_path}")

app = Flask(__name__)
os.environ["CUDA_DEVICE_ORDER"] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

from local_encoder import load_model, embedding

tokenizer, model = load_model(f"{os.getcwd()}/models/bge-large-zh-v1.5/")

@app.route('/embedding/', methods=["POST"])
def run_cls():
    resp_info = dict()
    if request.method == "POST":
        sentences = request.json.get('sentences')
        if sentences is not None and len(sentences) != 0:
            logger.info(sentences)
            rst = embedding(tokenizer, model, sentences)
            resp_info["code"] = 200
            resp_info["data"] = rst
        else:
            resp_info["msg"] = "Input is None, please check !"
            resp_info["code"] = 406
    resp = make_response(json.dumps(resp_info))
    resp.status_code = 200
    return resp

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5163, debug=False)