LongArticleAlgServer/alg_app.py

"""
@author: luojunhui
"""
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)
embedding_manager = EmbeddingManager()

@app.before_serving
async def init():
    """
    初始化模型
    """
    await AsyncMySQL.init_pool()
    model = BertSimilarity(model_name_or_path="BAAI/bge-large-zh-v1.5")
    print("模型加载成功")
    embedding_manager.set_model(model)
    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()