Refactor PoetryGenerationService to streamline model initialization and improve error handling

app/services/poetry_generation.py

@@ -1,27 +1,27 @@
-from typing import Optional, Dict, List
+from typing import Optional, List
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
 import logging
 from functools import lru_cache
 import concurrent.futures
-from torch.cuda import empty_cache
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+model_name = "meta-llama/Llama-3.2-1B-Instruct"
+
 class ModelManager:
     _instance = None
-    _initialized = False
-    _model_name = "meta-llama/Llama-3.2-1B-Instruct"
     
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super().__new__(cls)
+            cls._initialized = False
         return cls._instance
     
     def __init__(self):
-            
+        if not ModelManager._initialized:
             # Initialize tokenizer and model
             self.tokenizer = AutoTokenizer.from_pretrained(model_name)
             self.tokenizer.pad_token = self.tokenizer.eos_token
@@ -30,9 +30,7 @@ class ModelManager:
                 torch_dtype=torch.float16,
                 device_map="auto"
             )
-            
             # Set model to evaluation mode and move to GPU
-            self.model = self.model.to(self.model.device)
             self.model.eval()
             ModelManager._initialized = True
     
@@ -41,8 +39,8 @@ class ModelManager:
             del self.model
             del self.tokenizer
             torch.cuda.empty_cache()
-        except:
-            pass
+        except Exception as e:
+            logger.error(f"Error during cleanup: {str(e)}")
 
 @lru_cache(maxsize=1)
 def get_hf_token() -> str:
@@ -54,35 +52,23 @@ def get_hf_token() -> str:
             "Please set your Hugging Face access token."
         )
     return token
-model_name = "meta-llama/Llama-3.2-1B-Instruct"
+
 class PoetryGenerationService:
     def __init__(self):
         # Get model manager instance
         model_manager = ModelManager()
         self.model = model_manager.model
         self.tokenizer = model_manager.tokenizer
-        self.cache = {}
+
     def preload_models(self):
         """Preload the models during application startup"""
         try:
-            # Initialize tokenizer and model
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-
-            self.model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16,
-                device_map="auto"
-            )
-
-            # Set model to evaluation mode and move to GPU
-            self.model = self.model.to(self.model.device)
-            self.model.eval()
-
+            _ = ModelManager()  # Ensure ModelManager singleton is initialized
             logger.info("Models preloaded successfully")
         except Exception as e:
             logger.error(f"Error preloading models: {str(e)}")
             raise
+
     def generate_poem(
         self,
         prompt: str,
@@ -119,7 +105,7 @@ class PoetryGenerationService:
         except Exception as e:
             raise Exception(f"Error generating poem: {str(e)}")
 
-    def generate_poems(self, prompts: list[str]) -> list[str]:
+    def generate_poems(self, prompts: List[str]) -> List[str]:
         with concurrent.futures.ThreadPoolExecutor() as executor:
             poems = list(executor.map(self.generate_poem, prompts))
-        return poems
+        return poems