Update chatNT.py

chatNT.py

@@ -694,6 +694,7 @@ class TorchMultiOmicsModel(PreTrainedModel):
             vocab_size - 1
         )
 
+        outs = {}
         if bio_token_ids is None:
             projected_bio_embeddings = None
         else:
@@ -708,14 +709,18 @@ class TorchMultiOmicsModel(PreTrainedModel):
 
 
                 # Project these embeddings
-                projected_bio_embeddings = [
-                    self.projection_model(
+                projected_bio_embeddings = []
+                print("(debug) remember to remove loop for projected")
+                for bio_seq_num, bio_embeddings in enumerate(bio_embeddings_list):
+                    proj, output = self.projection_model(
                         bio_token_ids=bio_token_ids[:, bio_seq_num],
                         bio_embeddings=bio_embeddings,
                         english_token_ids=projection_english_tokens_ids,
                     )
-                    for bio_seq_num, bio_embeddings in enumerate(bio_embeddings_list)
-                ]
+                    projected_bio_embeddings.append(proj)
+                    for key in output.keys():
+                        outs[f"{key}_{bio_seq_num}"] = output[key]
+                
                 projected_bio_embeddings = torch.stack(projected_bio_embeddings, dim=1)
 
         # decode        
@@ -724,7 +729,8 @@ class TorchMultiOmicsModel(PreTrainedModel):
             projected_bio_embeddings=projected_bio_embeddings,
         )
 
-        outs = {"logits": logits, "projected_bio_embeddings": projected_bio_embeddings}
+        outs["logits"] = logits
+        outs["projected_bio_embeddings"] = projected_bio_embeddings
         
         # Just for debugging
         print("(debug) remember to remove bio_embeddings storage")
@@ -1848,8 +1854,12 @@ class TorchMultiModalPerceiverResamplerProjection(nn.Module):
             english_token_ids (torch.Tensor):
                 Shape (batch_size, num_english_tokens)
         """
+        outs = {}
         projected_bio_embeddings = self.bio_projection(bio_embeddings)
+        print("(debug) remember to remove this projected_bio_embeddings out, and 'outs' output")
+        outs['projected_bio_embeddings'] = projected_bio_embeddings
         english_embeddings = self.token_embedding(english_token_ids)
+        outs['english_embeddings'] = english_embeddings
 
         bio_attention_mask = build_perceiver_padding_attention_mask(
             bio_token_ids, self.config.resampled_length, self.bio_pad_token_id
@@ -1865,7 +1875,7 @@ class TorchMultiModalPerceiverResamplerProjection(nn.Module):
             attention_mask_2=english_attention_mask,
         )["embeddings"]
 
-        return projected_embeddings
+        return projected_embeddings, outs
 
 
 def build_perceiver_padding_attention_mask(