Make CiteFromPrompt robust against tokenizer vocab size discrepancies

Signed-off-by: aerdem4 <ahmeterd4@gmail.com>

Author: aerdem4

logits_processor_zoo/transformers/cite_prompt.py

@@ -41,11 +41,12 @@ def __init__(self, tokenizer: PreTrainedTokenizer, boost_factor: float = 1.0, bo
         self.boost_eos = boost_eos
 
     def _process(self, input_ids: List[int], scores: torch.Tensor) -> torch.Tensor:
+        voc_size = scores.shape[1]
         for i in range(scores.shape[0]):
             tokens = set(self.prompt_token_ids[i])
             if self.boost_eos:
                 tokens.add(self.eos_token_id)
 
-            tokens = list(tokens)
+            tokens = [t for t in tokens if t < voc_size]
             scores[i, tokens] += self.boost_factor
         return scores

logits_processor_zoo/vllm/cite_prompt.py

@@ -46,6 +46,6 @@ def __call__(self, prompt_tokens_ids: List[int], past_token_ids: List[int], scor
         if self.boost_eos:
             tokens.add(self.eos_token_id)
 
-        tokens = list(tokens)
+        tokens = [t for t in tokens if t < scores.shape[0]]
         scores[tokens] += self.boost_factor
         return scores