abhisheksan commited on
Commit
1c1ca6d
·
1 Parent(s): f1b3987

Enhance lifespan management in FastAPI by initializing PoetryGenerationService and handling model preloading asynchronously

Browse files
Files changed (1) hide show
  1. main.py +15 -6
main.py CHANGED
@@ -1,3 +1,4 @@
 
1
  from contextlib import asynccontextmanager
2
  from fastapi import FastAPI
3
  from app.api.endpoints.poetry import router as poetry_router
@@ -37,13 +38,21 @@ def init_huggingface():
37
 
38
  @asynccontextmanager
39
  async def lifespan(app: FastAPI):
40
- # Startup
41
- poetry_service = PoetryGenerationService()
42
- await poetry_service.preload_models()
43
  init_huggingface()
44
- yield
45
- # Cleanup (if needed)
46
- # Add cleanup code here
 
 
 
 
 
 
 
 
 
 
47
 
48
  app = FastAPI(lifespan=lifespan)
49
  app.include_router(poetry_router, prefix="/api/v1/poetry")
 
1
+ import asyncio
2
  from contextlib import asynccontextmanager
3
  from fastapi import FastAPI
4
  from app.api.endpoints.poetry import router as poetry_router
 
38
 
39
  @asynccontextmanager
40
  async def lifespan(app: FastAPI):
41
+ # Initialize Hugging Face authentication first
 
 
42
  init_huggingface()
43
+
44
+ # Initialize poetry service and preload models
45
+ poetry_service = PoetryGenerationService()
46
+
47
+ # Only await if `preload_models` is asynchronous
48
+ if callable(getattr(poetry_service, "preload_models", None)):
49
+ result = poetry_service.preload_models()
50
+ if asyncio.iscoroutine(result):
51
+ await result
52
+ else:
53
+ result() # Call directly if synchronous
54
+
55
+ yield # Continue to application startup
56
 
57
  app = FastAPI(lifespan=lifespan)
58
  app.include_router(poetry_router, prefix="/api/v1/poetry")