Add detailed logging and improve error handling in model manager

Dockerfile

@@ -26,6 +26,8 @@ ENV PORT=7860
 ENV PATH="/home/user/.local/bin:${PATH}"
 ENV HF_HOME=/home/user/.cache/huggingface
 ENV TRANSFORMERS_CACHE=/home/user/.cache/huggingface
+# Set logging to stdout
+ENV LOG_FILE=/dev/stdout
 
 # Switch to non-root user
 USER user
@@ -46,5 +48,5 @@ COPY --chown=user:user . .
 # Expose port for Hugging Face Spaces
 EXPOSE 7860
 
-# Run the application
-CMD ["python", "-m", "uvicorn", "src.api:app", "--host", "0.0.0.0", "--port", "7860"]
+# Run the application with logging
+CMD ["python", "-u", "-m", "uvicorn", "src.api:app", "--host", "0.0.0.0", "--port", "7860", "--log-level", "debug"]

src/model_manager.py

@@ -1,5 +1,5 @@
 import logging
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import torch
 from huggingface_hub import login
 from .config import Config
@@ -16,7 +16,12 @@ class ModelManager:
         # Login to Hugging Face Hub
         if Config.HUGGING_FACE_TOKEN:
             logger.info("Logging in to Hugging Face Hub")
-            login(token=Config.HUGGING_FACE_TOKEN)
+            try:
+                login(token=Config.HUGGING_FACE_TOKEN)
+                logger.info("Successfully logged in to Hugging Face Hub")
+            except Exception as e:
+                logger.error(f"Failed to login to Hugging Face Hub: {str(e)}")
+                raise
         
         # Initialize tokenizer and model
         self._init_tokenizer()
@@ -37,7 +42,8 @@ class ModelManager:
                 'bos_token': '<s>'
             }
             self.tokenizer.add_special_tokens(special_tokens)
-            logger.info("Tokenizer loaded successfully.")
+            logger.info("Tokenizer loaded successfully")
+            logger.debug(f"Tokenizer vocabulary size: {len(self.tokenizer)}")
         except Exception as e:
             logger.error(f"Error loading tokenizer: {str(e)}")
             raise
@@ -46,6 +52,7 @@ class ModelManager:
         """Initialize the model."""
         try:
             logger.info(f"Loading model: {self.model_name}")
+            logger.info(f"Using device: {self.device}")
             
             # Load model with CPU configuration
             self.model = AutoModelForCausalLM.from_pretrained(
@@ -57,7 +64,8 @@ class ModelManager:
             )
             # Resize embeddings to match tokenizer
             self.model.resize_token_embeddings(len(self.tokenizer))
-            logger.info(f"Using device: {self.device}")
+            logger.info("Model loaded successfully")
+            logger.debug(f"Model parameters: {sum(p.numel() for p in self.model.parameters())}")
         except Exception as e:
             logger.error(f"Error loading model: {str(e)}")
             raise
@@ -65,11 +73,16 @@ class ModelManager:
     def generate_text(self, prompt: str, max_new_tokens: int = 1024) -> str:
         """Generate text from prompt."""
         try:
+            logger.info("Starting text generation")
+            logger.debug(f"Prompt length: {len(prompt)}")
+            
             # Encode the prompt
             inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
             inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            logger.debug(f"Input tensor shape: {inputs['input_ids'].shape}")
 
             # Generate response
+            logger.info("Generating response")
             with torch.no_grad():
                 outputs = self.model.generate(
                     **inputs,
@@ -80,27 +93,16 @@ class ModelManager:
                     pad_token_id=self.tokenizer.pad_token_id,
                     eos_token_id=self.tokenizer.eos_token_id,
                 )
-
+            
             # Decode and return the generated text
             generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            # Extract only the generated part (remove the prompt)
             response = generated_text[len(prompt):].strip()
             
+            logger.info("Text generation completed")
+            logger.debug(f"Response length: {len(response)}")
             return response
             
         except Exception as e:
             logger.error(f"Error generating text: {str(e)}")
-            return """- Issues:
-- Error generating code review
-- Model inference failed
-
-- Improvements:
-- Please try again
-- Check model configuration
-
-- Best Practices:
-- Ensure proper model setup
-- Verify token permissions
-
-- Security:
-- No immediate concerns"""
+            logger.error(f"Error details: {type(e).__name__}")
+            raise