如何在aiogram 3启动时初始化ML模型并在处理器中调用？

我正在使用ML模型，需要在程序启动时初始化模型，并将模型与tokenizer传递给处理器以运行模型，该如何实现？目前我用FSMContext，每个用户发送/start命令时都会单独加载并保存模型，这样太浪费资源了。

当前项目结构及代码如下：

model.py #

def load_model():
    return model
def load_tokenizer():
    return tokenizer
def run_model(model, tokenizer):
    <...>
    return output

bot.py #

class TelegramBot:
    def __init__(self):
        self.TOKEN = "TOKEN"

        self.storage = MemoryStorage()
        self.dp = Dispatcher(storage=self.storage)

    async def main(self):
        self.bot = Bot(token=self.TOKEN)
        self.dp.include_routers(
            start_handler.router,
            main_handler.router,
        )
        await self.dp.start_polling(self.bot)
asyncio.run(TelegramBot().main())

main_handler.py #

async def main_handler(message: Message, state: FSMContext):
    data = await state.get_data()
    model = data["model"]
    tokenizer = data["tokenizer"]
    result = run_model(model, tokenizer)
    await message.answer(result)

start_handler.py #

from model import load_model, load_tokenizer

async def start_handler(message: Message, state: FSMContext):
    await state.set_data(
        {
            "model": load_model(),
            "tokenizer": load_tokenizer()
        }
    )

核心思路是在程序启动时只加载一次模型和tokenizer，作为全局实例注入到Dispatcher上下文，所有处理器共享这一个实例，避免重复加载消耗内存。

步骤1：优化model.py的模型加载逻辑 #

完善加载逻辑并实现全局单例加载（以HuggingFace Transformers为例）：

from transformers import AutoModel, AutoTokenizer

# 全局实例，程序启动时仅加载一次
MODEL = None
TOKENIZER = None

def load_model(model_name="your-model-name"):
    global MODEL
    if MODEL is None:
        MODEL = AutoModel.from_pretrained(model_name)
        MODEL.eval()  # 切换到推理模式
    return MODEL

def load_tokenizer(model_name="your-model-name"):
    global TOKENIZER
    if TOKENIZER is None:
        TOKENIZER = AutoTokenizer.from_pretrained(model_name)
    return TOKENIZER

def run_model(model, tokenizer, input_text):
    # 示例推理逻辑，根据实际模型调整
    inputs = tokenizer(input_text, return_tensors="pt")
    outputs = model(**inputs)
    return tokenizer.decode(outputs.logits.argmax(dim=-1)[0])

步骤2：修改bot.py，启动时预加载模型 #

在Bot初始化阶段加载模型和tokenizer，存入Dispatcher的全局数据池：

import asyncio
from aiogram import Bot, Dispatcher
from aiogram.fsm.storage.memory import MemoryStorage
import start_handler
import main_handler
from model import load_model, load_tokenizer

class TelegramBot:
    def __init__(self):
        self.TOKEN = "TOKEN"
        self.storage = MemoryStorage()
        self.dp = Dispatcher(storage=self.storage)
        # 程序启动时加载模型，存入Dispatcher全局数据
        self.dp["model"] = load_model()
        self.dp["tokenizer"] = load_tokenizer()

    async def main(self):
        self.bot = Bot(token=self.TOKEN)
        self.dp.include_routers(
            start_handler.router,
            main_handler.router,
        )
        await self.dp.start_polling(self.bot)

if __name__ == "__main__":
    asyncio.run(TelegramBot().main())

步骤3：修改start_handler.py，移除模型加载逻辑 #

/start命令仅处理用户状态初始化，不再加载模型：

from aiogram import Router
from aiogram.types import Message
from aiogram.fsm.context import FSMContext

router = Router()

@router.message(commands=["start"])
async def start_handler(message: Message, state: FSMContext):
    await state.set_data({"user_init": True})
    await message.answer("机器人已启动，可发送消息进行推理！")

步骤4：修改main_handler.py，从Dispatcher获取全局模型 #

直接从Dispatcher全局数据中读取模型和tokenizer：

from aiogram import Router
from aiogram.types import Message
from aiogram.fsm.context import FSMContext
from model import run_model

router = Router()

@router.message()
async def main_handler(message: Message, state: FSMContext, dp: Dispatcher):
    # 从Dispatcher全局数据获取模型实例
    model = dp["model"]
    tokenizer = dp["tokenizer"]
    # 执行推理并返回结果
    result = run_model(model, tokenizer, message.text)
    await message.answer(result)

关键说明 #

• 全局单例加载：彻底避免重复加载模型，大幅降低内存占用和启动延迟，大模型场景效果尤为明显。
• Dispatcher全局数据：这是Aiogram官方推荐的全局共享数据存储方式，所有处理器可通过参数直接访问。
• 异步优化：如果模型推理是同步阻塞操作，可使用asyncio.to_thread包装run_model，避免阻塞Bot事件循环。

内容的提问来源于stack exchange，提问作者Stepan