LongArticleAlgServer/alg_app.py

"""
@author: luojunhui
"""
import os
from quart import Quart
from similarities import BertSimilarity
from routes import AlgRoutes
from applications import AsyncMySQLClient
from applications.embedding_manager import EmbeddingManager

app = Quart(__name__)
AsyncMySQL = AsyncMySQLClient(app)


@app.before_serving
async def init():
    """
    初始化模型
    """
    await AsyncMySQL.init_pool()

    model = BertSimilarity(model_name_or_path="BAAI/bge-large-zh-v1.5", device="cuda:1")
    embedding_manager = EmbeddingManager(model)
    print("模型加载成功")
    app_routes = AlgRoutes(AsyncMySQL, model, embedding_manager)
    app.register_blueprint(app_routes)


@app.after_serving
async def close_db():
    """
    关闭连接
    :return:
    """
    await AsyncMySQL.close_pool()


if __name__ == '__main__':
    app.run()