Optimize model loading and error handling in PoetryGenerationService; implement async poem generation and enhance application startup process

.gitignore

@@ -0,0 +1,3 @@
+.env 
+__pycache__/
+__pycache__/main.cpython-312.pyc

app/services/poetry_generation.py

@@ -1,5 +1,7 @@
+# poetry_generation.py
+import asyncio
 from typing import Optional, List
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import torch
 import os
 import logging
@@ -9,7 +11,8 @@ import concurrent.futures
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-model_name = "meta-llama/Llama-3.2-1B-Instruct"
+# Use a smaller model
+model_name = "facebook/opt-125m"  # Much smaller than Llama
 
 class ModelManager:
     _instance = None
@@ -22,16 +25,35 @@ class ModelManager:
     
     def __init__(self):
         if not ModelManager._initialized:
-            # Initialize tokenizer and model
+            # Initialize quantization config
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16
+            )
+            
+            # Initialize tokenizer and model with quantization
             self.tokenizer = AutoTokenizer.from_pretrained(model_name)
             self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+            # Load model with optimizations
             self.model = AutoModelForCausalLM.from_pretrained(
                 model_name,
-                torch_dtype=torch.float16,
-                device_map="auto"
+                quantization_config=quantization_config,
+                device_map="auto",
+                torch_dtype=torch.float16
             )
-            # Set model to evaluation mode and move to GPU
+            
+            # Enable model optimizations
+            self.model.config.use_cache = True
+            
+            # Set model to evaluation mode
             self.model.eval()
+            
+            # Move model to GPU if available
+            if torch.cuda.is_available():
+                self.model = self.model.cuda()
+                torch.backends.cudnn.benchmark = True
+            
             ModelManager._initialized = True
     
     def __del__(self):
@@ -42,30 +64,23 @@ class ModelManager:
         except Exception as e:
             logger.error(f"Error during cleanup: {str(e)}")
 
-@lru_cache(maxsize=1)
-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
-
 class PoetryGenerationService:
     def __init__(self):
-        # Get model manager instance
         model_manager = ModelManager()
         self.model = model_manager.model
         self.tokenizer = model_manager.tokenizer
+        
+        # Pre-compile common prompt templates
+        self.prompt_template = "Write a short poem about {}\n"
 
     def preload_models(self):
         """Preload the models during application startup"""
         try:
-            _ = ModelManager()  # Ensure ModelManager singleton is initialized
+            _ = ModelManager()
+            # Warmup generation
+            self.generate_poem("warmup")
             logger.info("Models preloaded successfully")
-            return True  # Return a meaningful value
+            return True
         except Exception as e:
             logger.error(f"Error preloading models: {str(e)}")
             raise Exception("Failed to preload models") from e
@@ -76,14 +91,25 @@ class PoetryGenerationService:
         temperature: Optional[float] = 0.7,
         top_p: Optional[float] = 0.9,
         top_k: Optional[int] = 50,
-        max_length: Optional[int] = 100,
+        max_length: Optional[int] = 150,
         repetition_penalty: Optional[float] = 1.1
     ) -> str:
         try:
-            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=128)
+            # Format prompt using template
+            formatted_prompt = self.prompt_template.format(prompt)
+            
+            # Optimize input processing
+            inputs = self.tokenizer(
+                formatted_prompt,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=64  # Reduced from 128
+            )
+            
             inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
             
-            with torch.no_grad():
+            with torch.inference_mode():  # Faster than torch.no_grad()
                 outputs = self.model.generate(
                     inputs["input_ids"],
                     attention_mask=inputs["attention_mask"],
@@ -95,6 +121,8 @@ class PoetryGenerationService:
                     repetition_penalty=repetition_penalty,
                     pad_token_id=self.tokenizer.eos_token_id,
                     eos_token_id=self.tokenizer.eos_token_id,
+                    num_beams=1,  # Disable beam search for speed
+                    early_stopping=True
                 )
             
             return self.tokenizer.decode(
@@ -104,9 +132,14 @@ class PoetryGenerationService:
             )
             
         except Exception as e:
-            raise Exception(f"Error generating poem: {str(e)}")
+            logger.error(f"Error generating poem: {str(e)}")
+            return f"Error generating poem: {str(e)}"
 
-    def generate_poems(self, prompts: List[str]) -> List[str]:
+    async def generate_poems_async(self, prompts: List[str]) -> List[str]:
+        loop = asyncio.get_event_loop()
         with concurrent.futures.ThreadPoolExecutor() as executor:
-            poems = list(executor.map(self.generate_poem, prompts))
-        return poems
+            poems = await asyncio.gather(
+                *[loop.run_in_executor(executor, self.generate_poem, prompt) 
+                  for prompt in prompts]
+            )
+        return poems

main.py

@@ -1,6 +1,7 @@
+# main.py
 import asyncio
 from contextlib import asynccontextmanager
-from fastapi import FastAPI
+from fastapi import FastAPI, BackgroundTasks
 from app.api.endpoints.poetry import router as poetry_router
 import os
 import logging
@@ -11,10 +12,11 @@ from huggingface_hub import login
 from functools import lru_cache
 from app.services.poetry_generation import PoetryGenerationService
 
-# Configure logging once at module level
 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."""
@@ -35,26 +37,29 @@ def init_huggingface():
     except Exception as e:
         logger.error(f"Failed to login to Hugging Face: {str(e)}")
         raise
-
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Initialize Hugging Face authentication first
+    global poetry_service
+    
+    # Initialize Hugging Face authentication
     init_huggingface()
-
+    
     # Initialize poetry service and preload models
     poetry_service = PoetryGenerationService()
-
+    
     try:
-        preload_result = poetry_service.preload_models()
-        if asyncio.iscoroutine(preload_result):
-            await preload_result
-        else:
-            preload_result  # Call directly if synchronous
+        # 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 preloading models: {str(e)}")
+        logger.error(f"Error during startup: {str(e)}")
         raise
-
-    yield  # Continue to application startup
+    finally:
+        logger.info("Shutting down application")
 
 app = FastAPI(lifespan=lifespan)
 app.include_router(poetry_router, prefix="/api/v1/poetry")
@@ -63,14 +68,16 @@ app.include_router(poetry_router, prefix="/api/v1/poetry")
 async def lifecheck():
     return Response("OK", media_type="text/plain")
 
-def get_port() -> int:
-    return int(os.getenv("PORT", "8000"))
-
 if __name__ == "__main__":
     import uvicorn
     
-    port = get_port()
-    app.mount("/static", StaticFiles(directory="static"), name="static")
+    port = int(os.getenv("PORT", "8000"))
     
-    logger.info(f"Starting FastAPI server on port {port}")
-    uvicorn.run(app, host="0.0.0.0", port=port)
+    # 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
+    )