Update chatNT.py

chatNT.py

@@ -405,9 +405,7 @@ class TorchBioBrainDecoder(nn.Module):
         """
 
         # Compute English token embeddings
-        print("(debug) in biobraindecoder, english tokens ids : ", english_token_ids.shape)
         tokens_embeddings = self.gpt_model.token_embed(english_token_ids)
-        print("(debug) tokens_embeddings shape : ", tokens_embeddings.shape)
 
         if projected_bio_embeddings is not None:
             (
@@ -419,8 +417,10 @@ class TorchBioBrainDecoder(nn.Module):
             
             # Insert the bio embeddings at the SEQ token positions
             processed_tokens_ids = english_token_ids.clone()
-            print("(debug) Inside : processed tokens ids shape : ", processed_tokens_ids.shape)
-            print("(debug) Inside : projected bio embeddings shape : ", projected_bio_embeddings.shape)
+            print("(debug) Before call tokens embeddings shape : ", tokens_embeddings.shape)
+            print("(debug) Before call Processed tokens ids shape : ", processed_tokens_ids.shape)
+            print("(debug) Before call Projected bio embeddings shape : ", projected_bio_embeddings.shape)
+            print("num bio sequences : ", num_bio_sequences)
             for bio_seq_num in range(num_bio_sequences):
                 tokens_embeddings, processed_tokens_ids = self.insert_embeddings(
                     processed_tokens_ids,
@@ -431,7 +431,6 @@ class TorchBioBrainDecoder(nn.Module):
                 print("After call : ", tokens_embeddings.shape)
 
         # Regular GPT pass through
-        print("(debug) tokens embeddings shape : ", tokens_embeddings.shape)
         embeddings = self.gpt_model.apply_transformer_layers(tokens_embeddings)
         embeddings = self.gpt_model.final_norm(embeddings)
 
@@ -472,6 +471,11 @@ 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)
 
         def _insert(
             tokens_1d: torch.Tensor,
@@ -485,6 +489,7 @@ 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:
                 idx = indices[0].item()
@@ -501,6 +506,7 @@ class TorchBioBrainDecoder(nn.Module):
                     :-1, :
                 ]
                 tokens_1d[idx] = -1
+                print("_insert output : ", x.shape)
                 return x, tokens_1d
             else:
                 return (
@@ -519,8 +525,11 @@ 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
 
@@ -701,13 +710,11 @@ class TorchMultiOmicsModel(PreTrainedModel):
 
             if projected_bio_embeddings is None:
                 # Compute bio sequences embeddings
-                print("(debug) shape bio tokens ids : ", bio_token_ids.shape)
                 bio_embeddings_list = [
                     self.biobrain_encoder(bio_token_ids=bio_token_ids[:, bio_seq_num])
                     for bio_seq_num in range(num_bio_sequences)
                 ]
 
-                print("(debug) shape of embeddings : ", bio_embeddings_list[0].shape)
 
                 # Project these embeddings
                 projected_bio_embeddings = [
@@ -718,14 +725,9 @@ class TorchMultiOmicsModel(PreTrainedModel):
                     )
                     for bio_seq_num, bio_embeddings in enumerate(bio_embeddings_list)
                 ]
-                print("(debug) Shape output projection model : ", projected_bio_embeddings[0].shape)
                 projected_bio_embeddings = torch.stack(projected_bio_embeddings, dim=1)
-                print("(debug) Shape projected bio embeddings : ", projected_bio_embeddings.shape)
 
         # decode
-        print("(debug) Going in biobrain decoder : ")
-        print("(debug) English token ids : ", english_token_ids.shape)
-        print("(debug) Projected bio embeddings : ", projected_bio_embeddings.shape)
         logits = self.biobrain_decoder(
             english_token_ids=english_token_ids,
             projected_bio_embeddings=projected_bio_embeddings,
@@ -899,7 +901,6 @@ class TorchGptGroupedQueryAttention(nn.Module):
         value_inputs: torch.Tensor,
         attention_mask: torch.Tensor = None,
     ) -> torch.Tensor:
-        print("(debug) Query input shape : ", query_inputs.shape)
         batch_size, seq_len, _ = query_inputs.shape
 
         queries = self.query_linear(query_inputs).view(  # noqa
@@ -981,7 +982,6 @@ class TorchGptDecoder(nn.Module):
         if attention_mask is None:
             attention_mask = build_causal_attention_mask(1, embeddings.shape[1])
         for layer in self.layers:
-            print("Embedding shape in apply_transformer_layers : ", embeddings.shape)
             embeddings = layer(embeddings, attention_mask)
 
         return embeddings