Update harim_plus.py

harim_plus.py

@@ -232,10 +232,10 @@ class Harimplus_Scorer:
                                                     labels = tgt_in.input_ids.masked_fill(fill_ignore_mask, -100),
                                                     return_dict=True).logits
                 sent_lengths = tgt_mask.sum(-1)
-                ll_tok = self.log_likelihoods(s2s_logits, tgt_in.input_ids, 1)#tgt_mask)
+                ll_tok = self.log_likelihoods(s2s_logits, tgt_in.input_ids, tgt_mask)
                 ll = ll_tok.sum(-1) / sent_lengths
 
-                harim_tok = self.harim(s2s_logits, lm_logits, tgt_in.input_ids, 1)#tgt_mask)
+                harim_tok = self.harim(s2s_logits, lm_logits, tgt_in.input_ids, tgt_mask)
                 harim = harim_tok.sum(-1) / sent_lengths
 
                 harim_plus_normalized = (ll + self._lambda * harim) # loglikelihood + lambda * negative_harim (negative harim=-1* risk)