Update src/generation.py

src/generation.py

@@ -1,6 +1,8 @@
 import logging
+import os
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from typing import List, Dict
+import time
 
 logger = logging.getLogger(__name__)
 
@@ -14,6 +16,7 @@ class ResponseGenerator:
             cache_folder (str, optional): Directory to cache model files (default: None).
         """
         logger.info(f"Initializing ResponseGenerator with model: {model_name}, cache_folder: {cache_folder}")
+        start_time = time.time()
         try:
             # Log cache contents for debugging
             if cache_folder and os.path.exists(cache_folder):
@@ -24,11 +27,14 @@ class ResponseGenerator:
                 cache_dir=cache_folder,
                 local_files_only=True
             )
+            logger.info(f"Tokenizer loaded in {time.time() - start_time:.2f} seconds")
+            start_time = time.time()
             self.model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 cache_dir=cache_folder,
                 local_files_only=True
             )
+            logger.info(f"Model loaded in {time.time() - start_time:.2f} seconds")
         except Exception as e:
             logger.error(f"Failed to load transformer model: {str(e)}")
             raise