Update handler.py

handler.py

@@ -18,15 +18,18 @@ class EndpointHandler:
     def get_allowed_token_ids(self, vocab_list: List[str]) -> set[int]:
         """
         Generate a set of token IDs for a given list of allowed words.
-        Includes both plain and space-prefixed forms of each word.
+        Includes plain, space-prefixed, capitalized, and uppercase forms of each word.
         """
         allowed_ids = set()
         for word in vocab_list:
-            # Add plain and space-prefixed token IDs
-            for token_id in self.tokenizer.encode(word, add_special_tokens=False):
-                allowed_ids.add(token_id)
-            for token_id in self.tokenizer.encode(" " + word, add_special_tokens=False):
-                allowed_ids.add(token_id)
+            # Generate all variations: plain, space-prefixed, capitalized, and uppercase
+            variations = {word, " " + word, word.capitalize(), " " + word.capitalize(), word.upper(), " " + word.upper()}
+            
+            # Add token IDs for all variations
+            for variation in variations:
+                for token_id in self.tokenizer.encode(variation, add_special_tokens=False):
+                    allowed_ids.add(token_id)
+                    
         return allowed_ids
 
     def filter_allowed_tokens(self, input_ids: torch.Tensor, scores: np.ndarray, allowed_token_ids: set[int]) -> np.ndarray: