Spaces:
Running
Running
Optimize model loading and error handling in PoetryGenerationService; implement async poem generation and enhance application startup process
6dbb459
# 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 | |
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 | |
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") | |
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 | |
) |