small fixes

EMCarrami · Nov 9, 2023 · 9eab60c · 9eab60c
1 parent d4399da
commit 9eab60c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/cprt/model/cprt_model.py b/cprt/model/cprt_model.py
@@ -156,8 +156,8 @@ def on_validation_epoch_end(self) -> None:
         self.log_dict({f"metrics/val_{k}": v.mean() for k, v in rouge_scores.items()})
         self.val_rouge_scores.reset()
         for idx, layer in enumerate(self.cprt_llm.transformer.h):
-            self.log(f"gates/layer_{idx}_attn_gate", layer.attn_gate.item())
-            self.log(f"gates/layer_{idx}_ff_gate", layer.ff_gate.item())
+            self.log(f"gates/layer_{idx}_attn_gate", layer.cross_attn.attn_gate.item())
+            self.log(f"gates/layer_{idx}_ff_gate", layer.cross_attn.ff_gate.item())
         # log example outputs
         input_text = self.text_tokenizer.batch_decode(self.last_val_batch.info, skip_special_tokens=True)
         for in_txt, protein in zip(input_text, self.last_val_batch.protein):

diff --git a/cprt/model/helper_modules.py b/cprt/model/helper_modules.py
@@ -48,6 +48,7 @@ def forward(
         attention_mask: Optional[torch.FloatTensor] = None,
         head_mask: Optional[torch.FloatTensor] = None,
         encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
         use_cache: Optional[bool] = False,
         output_attentions: Optional[bool] = False,
     ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
@@ -73,7 +74,7 @@ def forward(
             attention_mask=attention_mask,
             head_mask=head_mask,
             encoder_hidden_states=None,
-            encoder_attention_mask=None,
+            encoder_attention_mask=encoder_attention_mask,
             use_cache=use_cache,
             output_attentions=output_attentions,
         )