Update bert_handler.py

bert_handler.py

@@ -38,7 +38,11 @@ class BERTHandler:
     
     def __del__(self):
         """Destructor to ensure cleanup when object is deleted"""
-        self._cleanup_model()
+        try:
+            self._cleanup_model()
+        except Exception:
+            # Ignore cleanup errors during shutdown
+            pass
     
     def _cleanup_model(self):
         """
@@ -48,33 +52,57 @@ class BERTHandler:
         if hasattr(self, 'model') and self.model is not None:
             print("🧹 Cleaning up existing model from VRAM...")
             
+            # Check if torch is still available (can be None during shutdown)
+            try:
+                import torch as torch_module
+                if torch_module is None:
+                    return
+            except (ImportError, AttributeError):
+                return
+            
             # Move model to CPU first to free GPU memory
-            if torch.cuda.is_available() and next(self.model.parameters(), None) is not None:
-                if next(self.model.parameters()).is_cuda:
-                    self.model = self.model.cpu()
+            try:
+                if torch_module.cuda.is_available() and next(self.model.parameters(), None) is not None:
+                    if next(self.model.parameters()).is_cuda:
+                        self.model = self.model.cpu()
+            except Exception:
+                # Continue cleanup even if moving to CPU fails
+                pass
             
             # Delete the model
-            del self.model
-            self.model = None
+            try:
+                del self.model
+                self.model = None
+            except Exception:
+                pass
             
             # Force garbage collection
-            gc.collect()
+            try:
+                gc.collect()
+            except Exception:
+                pass
             
             # Clear CUDA cache
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-                torch.cuda.synchronize()  # Ensure all CUDA operations complete
+            try:
+                if torch_module.cuda.is_available():
+                    torch_module.cuda.empty_cache()
+                    torch_module.cuda.synchronize()  # Ensure all CUDA operations complete
+            except Exception:
+                pass
             
             print("✅ Model cleanup complete")
     
     def _print_vram_usage(self, prefix=""):
         """Print current VRAM usage for monitoring"""
-        if torch.cuda.is_available():
-            allocated = torch.cuda.memory_allocated() / 1e9
-            reserved = torch.cuda.memory_reserved() / 1e9
-            print(f"🎯 {prefix}VRAM: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved")
-        else:
-            print(f"🎯 {prefix}CUDA not available")
+        try:
+            if torch.cuda.is_available():
+                allocated = torch.cuda.memory_allocated() / 1e9
+                reserved = torch.cuda.memory_reserved() / 1e9
+                print(f"🎯 {prefix}VRAM: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved")
+            else:
+                print(f"🎯 {prefix}CUDA not available")
+        except Exception:
+            print(f"🎯 {prefix}VRAM: Could not read CUDA memory")
     
     def load_fresh_model(self, model_name="nomic-ai/nomic-bert-2048"):
         """Load fresh model and add special tokens with proper VRAM management"""
@@ -152,7 +180,17 @@ class BERTHandler:
             print(f"   - Embedding size: {self.model.bert.embeddings.word_embeddings.weight.shape[0]}")
             print(f"   - Tokenizer size: {len(self.tokenizer)}")
             
-            # DO NOT MODIFY ANYTHING - checkpoint is self-consistent
+            # Check for vocab size mismatch and warn (but don't auto-fix for checkpoints)
+            tokenizer_size = len(self.tokenizer)
+            model_vocab_size = self.model.config.vocab_size
+            embedding_size = self.model.bert.embeddings.word_embeddings.weight.shape[0]
+            
+            if not (tokenizer_size == model_vocab_size == embedding_size):
+                print(f"⚠️  VOCAB SIZE MISMATCH DETECTED:")
+                print(f"   - Tokenizer size: {tokenizer_size}")
+                print(f"   - Model config vocab_size: {model_vocab_size}")
+                print(f"   - Embedding size: {embedding_size}")
+                print(f"   This might affect inference quality.")
             
             # Load training state
             self._load_training_state(checkpoint_path)