Tokenization fixed

tokenization_nicheformer.py

@@ -257,27 +257,27 @@ class NicheformerTokenizer(PreTrainedTokenizer):
             # Avoid division by zero
             safe_mean = np.maximum(self.technology_mean, 1e-6)
             x = x / safe_mean
-        
+
         # Apply log1p transformation
         x = np.log1p(x)
-        
+
         # Convert to tokens
         tokens = _sub_tokenize_data(x, self.max_length, self.aux_tokens)
-        
+
         return tokens.astype(np.int32)
-    
+
     def __call__(self, data: Union[ad.AnnData, np.ndarray], **kwargs) -> Dict[str, torch.Tensor]:
         """Tokenize gene expression data.
-        
+
         Args:
             data: AnnData object or numpy array of gene expression data
-            
+
         Returns:
             Dictionary with input_ids and attention_mask tensors
         """
         if isinstance(data, ad.AnnData):
             adata = data.copy()
-            
+
             # Align with reference model if available
             if hasattr(self, '_load_reference_model'):
                 reference_model = self._load_reference_model()
@@ -285,15 +285,15 @@ class NicheformerTokenizer(PreTrainedTokenizer):
                     # Concatenate and then remove the reference
                     adata = ad.concat([reference_model, adata], join='outer', axis=0)
                     adata = adata[1:]
-            
+
             # Get gene expression data
             X = adata.X
-            
+
             # Get metadata for special tokens
             modality = adata.obs.get('modality', None)
             species = adata.obs.get('specie', None)
             technology = adata.obs.get('assay', None)
-            
+
             print(f"Modality: {modality}")
             print(f"Species: {species}")
             print(f"Technology: {technology}")
@@ -302,16 +302,18 @@ class NicheformerTokenizer(PreTrainedTokenizer):
                 modality_tokens = modality.astype(int).tolist()
             else:
                 modality_tokens = [self.modality_dict.get(m, self._vocabulary["[PAD]"]) for m in modality] if modality is not None else None
-            
+
             if species is not None and pd.api.types.is_numeric_dtype(species):
                 species_tokens = species.astype(int).tolist()
             else:
                 species_tokens = [self.species_dict.get(s, self._vocabulary["[PAD]"]) for s in species] if species is not None else None
-            
+
             if technology is not None and pd.api.types.is_numeric_dtype(technology):
                 technology_tokens = technology.astype(int).tolist()
+                print(f"Technology tokens: {technology_tokens}")
             else:
                 technology_tokens = [self.technology_dict.get(t, self._vocabulary["[PAD]"]) for t in technology] if technology is not None else None
+                print(f"Technology tokens resort: {technology_tokens}")
         else:
             X = data
             modality_tokens = None