Model: Fix logit bias handling

If the token doesn't exist, gracefully warn instead of erroring out.

Signed-off-by: kingbri <[email protected]>

backends/exllamav2/model.py

@@ -750,7 +750,13 @@ def generate_gen(self, prompt: str, **kwargs):
 
             # Map logits to the tensor with their biases
             for token, bias in logit_bias.items():
-                gen_settings.token_bias[token] = bias
+                if token in gen_settings.token_bias:
+                    gen_settings.token_bias[token] = bias
+                else:
+                    logger.warning(
+                        f"Logit bias: Token {token} not present "
+                        "in the model's vocab. Skipping."
+                    )
 
         # Ban the EOS token if specified. If not, append to stop conditions
         # as well.

common/sampling.py

@@ -111,7 +111,7 @@ class BaseSamplerRequest(BaseModel):
 
     logit_bias: Optional[Dict[int, float]] = Field(
         default_factory=lambda: get_default_sampler_value("logit_bias"),
-        examples=[[{"1": 10}]],
+        examples=[{"1": 10, "2": 50}],
     )
 
     negative_prompt: Optional[str] = Field(