bug fixed.

Author: msafari18

barcodebert/datasets.py

@@ -181,7 +181,18 @@ def __init__(
             if self.return_taxonomy_level:
                 taxonomy_column = f"{self.return_taxonomy_level}_name"
                 if taxonomy_column in df.columns:
-                    self.taxonomy_labels, self.taxonomy_label_set = pd.factorize(df[taxonomy_column], sort=True)
+                    # Replace empty strings and NaN with 'UNKNOWN'
+                    taxonomy_col = df[taxonomy_column].replace("", "UNKNOWN").fillna("UNKNOWN")
+                    self.taxonomy_labels, self.taxonomy_label_set = pd.factorize(taxonomy_col, sort=True)
+                    # Map 'UNKNOWN' samples to -1 so they're excluded from pair creation
+                    unknown_mask = taxonomy_col == "UNKNOWN"
+                    num_unknown = unknown_mask.sum()
+                    self.taxonomy_labels = [
+                        -1 if is_unknown else label for label, is_unknown in zip(self.taxonomy_labels, unknown_mask)
+                    ]
+                    print(f"Taxonomy labels: {len(self.taxonomy_labels)} total, {num_unknown} marked as UNKNOWN (-1)")
+                    print(f"Unique taxonomy categories: {self.taxonomy_label_set}")
+
                 else:
                     print(f"Warning: Column '{taxonomy_column}' not found. Using dummy labels.")
                     self.taxonomy_labels = [0] * len(self.labels)
@@ -195,7 +206,17 @@ def __init__(
             if self.return_taxonomy_level:
                 taxonomy_column = f"{self.return_taxonomy_level}_name"
                 if taxonomy_column in df.columns:
-                    self.taxonomy_labels, self.taxonomy_label_set = pd.factorize(df[taxonomy_column], sort=True)
+                    # Replace empty strings and NaN with 'UNKNOWN'
+                    taxonomy_col = df[taxonomy_column].replace("", "UNKNOWN").fillna("UNKNOWN")
+                    self.taxonomy_labels, self.taxonomy_label_set = pd.factorize(taxonomy_col, sort=True)
+                    # Map 'UNKNOWN' samples to -1 so they're excluded from pair creation
+                    unknown_mask = taxonomy_col == "UNKNOWN"
+                    num_unknown = unknown_mask.sum()
+                    self.taxonomy_labels = [
+                        -1 if is_unknown else label for label, is_unknown in zip(self.taxonomy_labels, unknown_mask)
+                    ]
+                    print(f"Taxonomy labels: {len(self.taxonomy_labels)} total, {num_unknown} marked as UNKNOWN (-1)")
+                    print(f"Unique taxonomy categories: {self.taxonomy_label_set}")
                 else:
                     print(f"Warning: Column '{taxonomy_column}' not found. Using dummy labels.")
                     self.taxonomy_labels = [0] * len(self.labels)

barcodebert/jumbo_transformer.py

@@ -27,7 +27,7 @@ def __init__(self, embed_dim: int, jumbo_multiplier: int = 6, dropout: float = 0
 
         self.jumbo_mlp = nn.Sequential(
             nn.LayerNorm(self.jumbo_width),
-            nn.Linear(self.jumbo_width, self.jumbo_width * 2),  # Wide hidden layer X4
+            nn.Linear(self.jumbo_width, self.jumbo_width * 2),  # Wide hidden layer X2
             nn.GELU(),
             nn.Dropout(dropout),
             nn.Linear(self.jumbo_width * 2, self.jumbo_width),

barcodebert/jumbo_transformer_with_taxonomy.py

@@ -37,11 +37,8 @@ def __init__(
         if self.enable_taxonomy_classification:
             jumbo_dim = bert_config.hidden_size * jumbo_multiplier
             self.taxonomy_classifier = JumboTaxonomyClassifier(jumbo_dim=jumbo_dim)
-            # Alias for backward compatibility
-            self.genus_classifier = self.taxonomy_classifier
         else:
             self.taxonomy_classifier = None
-            self.genus_classifier = None
 
     def forward(
         self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, inputs_embeds=None, **kwargs
@@ -88,6 +85,11 @@ def bert(self):
         """Return the BERT model from the underlying transformer for compatibility."""
         return self.transformer.bert
 
+    @property
+    def genus_classifier(self):
+        """Backward compatibility alias for taxonomy_classifier."""
+        return self.taxonomy_classifier
+
 
 def create_jumbo_transformer_with_taxonomy(
     bert_config, jumbo_multiplier=6, share_jumbo_mlp_across_layers=False, enable_taxonomy_classification=True

barcodebert/maelm_model.py

@@ -64,11 +64,8 @@ def __init__(
                 jumbo_dim = decoder_config.hidden_size * jumbo_multiplier
 
             self.taxonomy_classifier = JumboTaxonomyClassifier(jumbo_dim=jumbo_dim)
-            # Keep genus_classifier as alias for backward compatibility
-            self.genus_classifier = self.taxonomy_classifier
         else:
             self.taxonomy_classifier = None
-            self.genus_classifier = None
 
     def forward(self, input_ids, attention_mask, mask_positions, model_type="maelm_v2"):
         if model_type == "maelm_v2":
@@ -243,3 +240,8 @@ def forward_baseline(self, input_ids, attention_mask):
         )
 
         return outputs
+
+    @property
+    def genus_classifier(self):
+        """Backward compatibility alias for taxonomy_classifier."""
+        return self.taxonomy_classifier

barcodebert/pretraining.py

@@ -675,6 +675,14 @@ def print_pass(*args, **kwargs):
         print(f"  Throughput .........{train_stats['throughput']:11.2f} samples/sec")
         print(f"  Loss ...............{train_stats['loss']:14.5f}")
         print(f"  Accuracy ...........{train_stats['accuracy']:11.2f} %")
+
+        # Print taxonomy classification metrics if available
+        taxonomy_level_display = taxonomy_level.capitalize()
+        if f"{taxonomy_level}_loss" in train_stats:
+            print(f"  {taxonomy_level_display} Loss .........{train_stats[f'{taxonomy_level}_loss']:14.5f}")
+            print(f"  {taxonomy_level_display} Accuracy .....{train_stats[f'{taxonomy_level}_accuracy'] * 100:11.2f} %")
+            print(f"  {taxonomy_level_display} Pairs ........{train_stats[f'{taxonomy_level}_pairs']:8d}")
+
         print(flush=True)
 
         # Validate ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -716,6 +724,12 @@ def print_pass(*args, **kwargs):
         print(f"  Loss ...............{eval_stats['loss']:14.5f}")
         print(f"  Accuracy ...........{eval_stats['accuracy']:11.2f} %")
 
+        # Print taxonomy classification metrics if available
+        if f"{taxonomy_level}_loss" in eval_stats:
+            print(f"  {taxonomy_level_display} Loss .........{eval_stats[f'{taxonomy_level}_loss']:14.5f}")
+            print(f"  {taxonomy_level_display} Accuracy .....{eval_stats[f'{taxonomy_level}_accuracy'] * 100:11.2f} %")
+            print(f"  {taxonomy_level_display} Pairs ........{eval_stats[f'{taxonomy_level}_pairs']:8d}")
+
         # Save model ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
         t_start_save = time.time()
 
@@ -2092,7 +2106,7 @@ def get_parser():
         "--jumbo_source",
         dest="jumbo_source",
         type=str,
-        default="encoder",
+        default="decoder",
         choices=["encoder", "decoder"],
         help=(
             "Source of jumbo tokens for taxonomy classification: 'encoder' (direct from encoder) or "