Как лучше всего передавать модели машинного обучения в FastAPI на конечные точки на разных маршрутизаторах? ⇐ Python

[Anonymous]

У меня есть приложение FastAPI, и я хочу использовать одну и ту же модель машинного обучения на разных конечных точках (запросы, обработка ввода и т. д. различаются).
теперь в моем (корневом) файле app.py У меня есть что-то вроде:

Код: Выделить всё

from fastapi import FastAPI
from transformers import AutoModelForCausalLM, AutoTokenizer

from api.routers.subrouter1 import router as sr1
from api.routers.subrouter2 import router as sr2

app = FastAPI()

@app.on_event("startup")
def load_model(model_tag:str):
app.state.model = AutoModelForCausalLM.from_pretrained(model_tag)
app.state.tokenizer = AutoTokenizer.from_pretrained(model_tag)

app.include_router(sr1, prefix="/sr1")
app.include_router(sr2, prefix="/sr2")

Вся загрузка работает, но у меня возникают проблемы, когда я хочу получить доступ к моделям для выполнения каких-либо выводов.
файл Subrouter 1 выглядит примерно так:

Код: Выделить всё

from fastapi import APIrouter

from api.models.subrouter1 import sr1Request

@router.post("/inference/")
def do_some_inference(request: sr1Request):
model = request.app.state.model
tokenizer = request.app.state.tokenizer
results = some_ml_function(request, model, tokenizer)
return results

Это общее представление о том, как сейчас структурирован код. Раньше я загружал модель в app.py вне какой-либо функции и просто назначал ее глобальному объекту модели/токенизатора, но у меня также были там все мои конечные точки, чтобы они могли сразу же ссылаться на нее. Теперь, когда приложение расширяется, мне нужны аккуратные маленькие маршрутизаторы, каждый из которых содержит свою логику. Затем я прочитал, что вы можете получить доступ к состоянию приложения через объект Request Starlette, и это тоже сработало. Однако теперь я также хочу использовать проверку объекта pydantic.
Теперь, когда я попадаю в конечную точку, и она анализирует объект пользовательского запроса, содержащий все мои поля и валидаторы, он естественно, у меня больше нет доступа к атрибуту request.app.state.model, и я получаю сообщение об ошибке.
Как лучше всего загружать модели ML только один раз в FastAPI и все равно иметь возможность использовать собственные объекты запроса? Кажется, это довольно обычный вариант использования, и я уверен, что упускаю что-то очевидное.

Подробнее здесь: https://stackoverflow.com/questions/790 ... fferent-ro