Add attn bias arg to HF wrapper (#458)

allenai · Feb 22, 2024 · cc36709 · cc36709
1 parent 7f7abbb
commit cc36709
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/hf_olmo/modeling_olmo.py b/hf_olmo/modeling_olmo.py
@@ -50,6 +50,7 @@ def forward(
         input_ids: torch.LongTensor = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
+        attention_bias: Optional[torch.Tensor] = None,
         past_key_values: Optional[List[torch.FloatTensor]] = None,
         labels: Optional[torch.LongTensor] = None,
         use_cache: Optional[bool] = None,
@@ -70,6 +71,7 @@ def forward(
             input_ids=input_ids,
             input_embeddings=inputs_embeds,
             attention_mask=attention_mask,
+            attention_bias=attention_bias,
             past_key_values=past_key_values,
             use_cache=use_cache,
             output_hidden_states=output_hidden_states,