Update handler.py

handler.py

@@ -1,52 +1,45 @@
 import torch
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    pipeline,
-    LogitsProcessor,
-    LogitsProcessorList
-)
+from llama_cpp import Llama  # Library for GGUF model handling
 from typing import Any, List, Dict
 
 
-class FixedVocabLogitsProcessor(LogitsProcessor):
+class FixedVocabLogitsProcessor:
     """
-    A custom LogitsProcessor that restricts the vocabulary
-    to a fixed set of token IDs, masking out everything else.
+    A custom logits processor for GGUF-compatible models.
     """
 
     def __init__(self, allowed_ids: set[int], fill_value=float('-inf')):
-        """
-        Args:
-          allowed_ids (set[int]): Token IDs allowed for generation.
-          fill_value (float): Value used to mask disallowed tokens, default -inf.
-        """
         self.allowed_ids = allowed_ids
         self.fill_value = fill_value
 
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+    def apply(self, logits: torch.FloatTensor):
         """
-        Args:
-          input_ids: shape (batch_size, sequence_length)
-          scores: shape (batch_size, vocab_size) - pre-softmax logits for the next token
-        Returns:
-          scores: shape (batch_size, vocab_size) with masked logits
+        Modify logits to restrict to allowed token IDs.
         """
-        batch_size, vocab_size = scores.size()
-        for b in range(batch_size):
-            for token_id in range(vocab_size):
-                if token_id not in self.allowed_ids:
-                    scores[b, token_id] = self.fill_value
-        return scores
+        for token_id in range(len(logits)):
+            if token_id not in self.allowed_ids:
+                logits[token_id] = self.fill_value
+        return logits
 
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # Load tokenizer and model
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=torch.float16)
+        """
+        Initialize the GGUF model handler.
+        Args:
+            path (str): Path to the GGUF file.
+        """
+        self.model = Llama(model_path=path)
+        self.tokenizer = self.model.tokenizer  # GGUF-specific tokenizer, if available
 
     def __call__(self, data: Any) -> List[Dict[str, str]]:
+        """
+        Handle the request, performing inference with a restricted vocabulary.
+        Args:
+            data (Any): Input data.
+        Returns:
+            List[Dict[str, str]]: Generated output.
+        """
         # Extract inputs and parameters
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
@@ -58,29 +51,20 @@ class EndpointHandler:
         # Define allowed tokens dynamically
         allowed_ids = set()
         for word in vocab_list:
-            for tid in self.tokenizer.encode(word, add_special_tokens=False):
-                allowed_ids.add(tid)
-            for tid in self.tokenizer.encode(" " + word, add_special_tokens=False):
+            for tid in self.model.tokenize(word):
                 allowed_ids.add(tid)
 
-        # Create custom logits processor
-        logits_processors = LogitsProcessorList([FixedVocabLogitsProcessor(allowed_ids=allowed_ids)])
-
-        # Prepare input IDs
-        input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids.to(self.model.device)
+        # Tokenize input
+        input_ids = self.model.tokenize(inputs)
 
-        # Generate output
+        # Perform inference
         output_ids = self.model.generate(
-            input_ids=input_ids,
-            logits_processor=logits_processors,
-            max_length=parameters.get("max_length", 30),
-            num_beams=parameters.get("num_beams", 1),
-            do_sample=parameters.get("do_sample", False),
-            pad_token_id=self.tokenizer.eos_token_id,
-            no_repeat_ngram_size=parameters.get("no_repeat_ngram_size", 3)
+            input_ids,
+            max_tokens=parameters.get("max_length", 30),
+            logits_processor=lambda logits: FixedVocabLogitsProcessor(allowed_ids).apply(logits)
         )
 
         # Decode the output
-        generated_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        generated_text = self.model.detokenize(output_ids)
 
         return [{"generated_text": generated_text}]