Update chatNT.py

chatNT.py

@@ -471,11 +471,8 @@ class TorchBioBrainDecoder(nn.Module):
                 - input_embeddings with resampled_embeddings inserted at the SEQ token
                 - tokens with the SEQ token set to -1
         """
-        print("Insert_embeddings input shape : ")
-        print("Tokens : ", tokens.shape)
-        print("Input embeddings : ", input_embeddings.shape)
-        print("Resampled embeddings : ", resampled_embeddings.shape)
-        print("Bio seq num : ", bio_seq_num)
+        print("Tokens : ", list(tokens))
+        print("seq_token_id : ", self.seq_token_id)
 
         def _insert(
             tokens_1d: torch.Tensor,
@@ -489,9 +486,9 @@ class TorchBioBrainDecoder(nn.Module):
                 resampled_embeddings (torch.Tensor):
                     Shape (bio_sequence_length, embed_dim,)
             """
-            print("_insert input : ", input_embeddings_1d.shape, resampled_embeddings_1d.shape)
             indices = torch.where(tokens_1d == self.seq_token_id)[0]
             if indices.numel() > 0:
+                print("going in if")
                 idx = indices[0].item()
                 insertion_pos = idx + resampled_embeddings_1d.shape[-2] * bio_seq_num
                 x = torch.cat(
@@ -506,9 +503,9 @@ class TorchBioBrainDecoder(nn.Module):
                     :-1, :
                 ]
                 tokens_1d[idx] = -1
-                print("_insert output : ", x.shape)
                 return x, tokens_1d
             else:
+                print("going in else")
                 return (
                     input_embeddings,
                     tokens_1d,
@@ -526,10 +523,8 @@ class TorchBioBrainDecoder(nn.Module):
             tokens_acc.append(tokens_out)
             embeddings_acc.append(embeddings_out)
 
-        print("(Embeddings_acc[0] shape : ", embeddings_acc[0].shape)
         tokens_acc = torch.stack(tokens_acc)
         embeddings_acc = torch.stack(embeddings_acc)
-        print("Embeddings acc shape : ", embeddings_acc.shape)
 
         return embeddings_acc, tokens_acc
 
@@ -703,6 +698,8 @@ class TorchMultiOmicsModel(PreTrainedModel):
             vocab_size - 1
         )
 
+        print("seq token id : ", self.seq_token_id)
+        print("Tokens at step 1 in multiomics : ", list(english_token_ids))
         if bio_token_ids is None:
             projected_bio_embeddings = None
         else:
@@ -728,6 +725,8 @@ class TorchMultiOmicsModel(PreTrainedModel):
                 projected_bio_embeddings = torch.stack(projected_bio_embeddings, dim=1)
 
         # decode
+        print("Tokens at step 2 in multiomics : ", list(english_token_ids))
+        
         logits = self.biobrain_decoder(
             english_token_ids=english_token_ids,
             projected_bio_embeddings=projected_bio_embeddings,