Change argument name in logits_bias

outlines/generate/generator.py

@@ -266,7 +266,7 @@ def expand_attention_masks(attention_masks: torch.Tensor) -> torch.Tensor:
 @torch.inference_mode()
 def bias_logits(
     logits: torch.Tensor,
-    ids_to_mask: List,
+    allowed_token_ids: List,
 ) -> torch.Tensor:
     """Mask the logits.
 
@@ -278,15 +278,15 @@ def bias_logits(
     logits
         Two dimensional tensor that contains the next-token probability
         distribution.
-    ids_to_mask
-        The ids to mask in each dimension.
+    allowed_token_ids
+        A list that contains the tokens that can be generated by the model.
 
     Returns
     -------
     A view of the original logits tensor where some values are masked.
 
     """
     biased_logits = torch.full(logits.shape, -math.inf, device=logits.device)
-    for i, ids in enumerate(ids_to_mask):
+    for i, ids in enumerate(allowed_token_ids):
         biased_logits[i, ids] = logits[i, ids]
     return biased_logits