first commit
This commit is contained in:
commit
5163bf819c
|
@ -0,0 +1,36 @@
|
|||
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)
|
Loading…
Reference in New Issue