Use a ml model loaded in the LIFESPAN function in the main file, in another file (ie: api router)

[pelaezluis]

First Check

• I added a very descriptive title here.
• I used the GitHub search to find a similar question and didn't find it.
• I searched the FastAPI documentation, with the integrated search.
• I already searched in Google "How to X in FastAPI" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to FastAPI but to Pydantic.
• I already checked if it is not related to FastAPI but to Swagger UI.
• I already checked if it is not related to FastAPI but to ReDoc.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from fastapi import FastAPI
from app.api import api
from contextlib import asynccontextmanager
import whisper

MODEL = 'base'
model = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    print(f'Loading {MODEL} model')
    model['whisper'] = whisper.load_model(MODEL)
    yield
    print('Shutting down the model...')

app = FastAPI(lifespan=lifespan)

@app.get('/')
async def root():
    return {'message': 'Server is running'}

app.include_router(api.router)

Description

I need to load a ML model, I understood the lifespan explained in documentation , but it is in the same file where are the endpoint, i need to load this model, but i have the api routes in different files, how can I use this loaded model?

Operating System

Windows

Operating System Details

No response

FastAPI Version

0.93.0

Python Version

3.10

Additional Context

No response

[kthy]

For inspiration, here's what I do with my database connection. In main.py:

from contextlib import asynccontextmanager
from fastapi import FastAPI
from .my_database_module import db_connect, db_disconnect


@asynccontextmanager
async def lifespan(_: FastAPI):
    await db_connect()
    yield
    await db_disconnect()


app = FastAPI(lifespan=lifespan)

Then, in my_database_module.py:

from databases import Database

DATABASE = Database("postgresql+asyncpg:my-username:hunter2@localhost:5432/db_name")


async def db_connect() -> None:
    await DATABASE.connect()


async def db_disconnect() -> None:
    if DATABASE.is_connected:
        await DATABASE.disconnect()

I could now import the DATABASE symbol in other files and use it there if I want to, but you should probably implement the methods that need to manipulate it in the imported module to avoid hard-coupling your code.

[pelaezluis]

Thanks! @kthy I'm gonna try this in my project

[jonra1993]

Hello @pelaezluis You can use a global context like this to share your model with other routes.

This sample code can give you an idea

main.py (It is not completed)

from app.utils.fastapi_globals import g, GlobalsMiddleware
from transformers import pipeline

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    # Load a pre-trained sentiment analysis model
    sentiment_model = pipeline("sentiment-analysis")
    g.set_default("sentiment_model", sentiment_model)
    print("startup fastapi")
    yield
    del sentiment_model
    g.cleanup()
    
app = FastAPI(
    title="Fastapi",
    lifespan=lifespan,
)

app.add_middleware(GlobalsMiddleware)

natural_language.py

from fastapi import APIRouter
from app.utils.fastapi_globals import g
from app.schemas.response_schema import IPostResponseBase, create_response

router = APIRouter()


@router.post("/sentiment_analysis")
async def sentiment_analysis_prediction(
    prompt: str = "Fastapi is awesome",
) -> IPostResponseBase:
    """
    Gets a sentimental analysis predition using a NLP model from transformers libray
    """
    sentiment_model = g.sentiment_model
    prediction = sentiment_model(prompt)
    return create_response(message="Prediction got succesfully", data=prediction)

A better example can be found this these two files main.py and natural_language.py

[pelaezluis]

This really help me! thanks a lot @jonra1993

[Arnold-git]

This was really helpful. Thanks for sharing!