Enhance ModelManager and PoetryGenerationService with optimizations and new features

- Update model name in ModelManager for improved performance
- Integrate optimized module for memory efficiency
- Add caching for Hugging Face token retrieval
- Modify generate_poem method to include truncation and max_length adjustments
- Introduce generate_poems method for batch poem generation using threading

app/services/poetry_generation.py

@@ -1,6 +1,18 @@
 from typing import Optional
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import os
+import logging
+from functools import lru_cache
+import concurrent.futures
+from torch.compile import (
+    InputsAreOptimized,
+    optimized_module,
+    optimized_static_function,
+)
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class ModelManager:
     _instance = None
@@ -13,7 +25,7 @@ class ModelManager:
     
     def __init__(self):
         if not ModelManager._initialized:
-            model_name = "meta-llama/Llama-3.2-3B-Instruct"
+            model_name = "meta-llama/Llama-2B-Instruct"
             
             # Initialize tokenizer and model
             self.tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -24,8 +36,10 @@ class ModelManager:
                 device_map="auto"
             )
             
-            # Set model to evaluation mode
+            # Set model to evaluation mode and move to GPU
+            self.model = self.model.to(self.model.device)
             self.model.eval()
+            self.model = optimized_module(self.model, memory_efficient=True)
             ModelManager._initialized = True
     
     def __del__(self):
@@ -36,24 +50,37 @@ class ModelManager:
         except:
             pass
 
+@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
+        self.cache = {}
     
-    async def generate_poem(
+    @optimized_static_function(InputsAreOptimized())
+    def generate_poem(
         self,
         prompt: str,
         temperature: Optional[float] = 0.7,
         top_p: Optional[float] = 0.9,
         top_k: Optional[int] = 50,
-        max_length: Optional[int] = 200,
+        max_length: Optional[int] = 100,
         repetition_penalty: Optional[float] = 1.1
     ) -> str:
         try:
-            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True)
+            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=128)
             inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
             
             with torch.no_grad():
@@ -77,4 +104,9 @@ class PoetryGenerationService:
             )
             
         except Exception as e:
-            raise Exception(f"Error generating poem: {str(e)}")
+            raise Exception(f"Error generating poem: {str(e)}")
+
+    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