first commit

This commit is contained in:
赵敬皓 2024-12-25 12:15:08 +08:00
commit 5163bf819c
1 changed files with 36 additions and 0 deletions

36
run.py Normal file
View File

@ -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)