Update handler.py

handler.py

@@ -2,7 +2,7 @@ from typing import Any, List, Dict
 from llama_cpp import Llama
 import numpy as np
 import torch
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, LogitsProcessorList
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -49,15 +49,16 @@ class EndpointHandler:
         if not inputs:
             raise ValueError("The 'inputs' field is required.")
 
-        logits_processor = None
-
+        # Prepare logits processor
+        logits_processors = None
         if vocab_list:
             # Define allowed tokens dynamically
             allowed_token_ids = self.get_allowed_token_ids(vocab_list)
 
-            # Define the logits processor if vocab_list is provided
-            def logits_processor(input_ids, scores):
-                return self.filter_allowed_tokens(input_ids, scores, allowed_token_ids)
+            # Create LogitsProcessorList with filtering function
+            logits_processors = LogitsProcessorList([
+                lambda input_ids, scores: self.filter_allowed_tokens(input_ids, scores, allowed_token_ids)
+            ])
 
         # Tokenize input
         input_ids = torch.tensor([self.tokenizer.encode(inputs, add_special_tokens=False)])
@@ -68,7 +69,7 @@ class EndpointHandler:
                 {"role": "user", "content": inputs}
             ],
             max_tokens=parameters.get("max_length", 30),
-            logits_processor=logits_processor,
+            logits_processor=logits_processors,  # Pass the LogitsProcessorList here
             temperature=parameters.get("temperature", 1),
             repeat_penalty=parameters.get("repeat_penalty", 1.0)
         )