poetica / main.py
abhisheksan's picture
Optimize model loading and error handling in PoetryGenerationService; implement async poem generation and enhance application startup process
6dbb459
raw
history blame
2.42 kB
# main.py
import asyncio
from contextlib import asynccontextmanager
from fastapi import FastAPI, BackgroundTasks
from app.api.endpoints.poetry import router as poetry_router
import os
import logging
from typing import Tuple
from starlette.responses import Response
from starlette.staticfiles import StaticFiles
from huggingface_hub import login
from functools import lru_cache
from app.services.poetry_generation import PoetryGenerationService
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Global poetry service instance
poetry_service = None
@lru_cache()
def get_hf_token() -> str:
"""Get Hugging Face token from environment variables."""
token = os.getenv("HF_TOKEN")
if not token:
raise EnvironmentError(
"HF_TOKEN environment variable not found. "
"Please set your Hugging Face access token."
)
return token
def init_huggingface():
"""Initialize Hugging Face authentication."""
try:
token = get_hf_token()
login(token=token)
logger.info("Successfully logged in to Hugging Face")
except Exception as e:
logger.error(f"Failed to login to Hugging Face: {str(e)}")
raise
@asynccontextmanager
async def lifespan(app: FastAPI):
global poetry_service
# Initialize Hugging Face authentication
init_huggingface()
# Initialize poetry service and preload models
poetry_service = PoetryGenerationService()
try:
# Preload models in background
background_tasks = BackgroundTasks()
background_tasks.add_task(poetry_service.preload_models)
logger.info("Application startup complete")
yield
except Exception as e:
logger.error(f"Error during startup: {str(e)}")
raise
finally:
logger.info("Shutting down application")
app = FastAPI(lifespan=lifespan)
app.include_router(poetry_router, prefix="/api/v1/poetry")
@app.get("/healthz")
async def lifecheck():
return Response("OK", media_type="text/plain")
if __name__ == "__main__":
import uvicorn
port = int(os.getenv("PORT", "8000"))
# Configure uvicorn with optimized settings
uvicorn.run(
app,
host="0.0.0.0",
port=port,
loop="uvloop", # Faster event loop implementation
http="httptools", # Faster HTTP protocol implementation
)