custom 2

handler.py

@@ -1,23 +1,86 @@
-from typing import Dict, List, Any
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    pipeline,
+    LogitsProcessor,
+    LogitsProcessorList
+)
+from typing import Any, List, Dict
+
+
+class FixedVocabLogitsProcessor(LogitsProcessor):
+    """
+    A custom LogitsProcessor that restricts the vocabulary
+    to a fixed set of token IDs, masking out everything else.
+    """
+
+    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:
+        """
+        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
+        """
+        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
 
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # load the model
+        # Load tokenizer and model
         tokenizer = AutoTokenizer.from_pretrained(path)
-        model = AutoModelForCausalLM.from_pretrained(path, device_map="auto")
-        # create inference pipeline
-        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=torch.float16)
+
+        # Define allowed tokens
+        words = ["Paris", "France", "Hello"]  # Customize as needed
+        allowed_ids = set()
+        for word in words:
+            for tid in tokenizer.encode(word, add_special_tokens=False):
+                allowed_ids.add(tid)
+            for tid in tokenizer.encode(" " + word, add_special_tokens=False):
+                allowed_ids.add(tid)
+
+        # Create custom logits processor
+        self.logits_processors = LogitsProcessorList([FixedVocabLogitsProcessor(allowed_ids=allowed_ids)])
 
-    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        self.tokenizer = tokenizer
+        self.model = model
+
+    def __call__(self, data: Any) -> List[Dict[str, str]]:
+        # Extract inputs and parameters
         inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", None)
-
-        # pass inputs with all kwargs in data
-        if parameters is not None:
-            prediction = self.pipeline(inputs, **parameters)
-        else:
-            prediction = self.pipeline(inputs)
-        # postprocess the prediction
-        return prediction
+        parameters = data.pop("parameters", {})
+
+        # Prepare input IDs
+        input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids.to(self.model.device)
+
+        # Generate output
+        output_ids = self.model.generate(
+            input_ids=input_ids,
+            logits_processor=self.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)
+        )
+
+        # Decode the output
+        generated_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+        return [{"generated_text": generated_text}]