Update handler

README.md

@@ -46,6 +46,7 @@ tokens = model.generate(
     eos_token_id=tokenizer.eos_token_id,
     early_stopping=True
 )
+print(tokenizer.decode(tokens[0]))
 
 # The hobbits were so suprised seeing their friend again that they did not 
 # speak. Aragorn looked at them, and then he turned to the others.</s>

handler.py

@@ -1,9 +1,13 @@
 from typing import Dict, List, Any
+import logging
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftConfig, PeftModel
 
 
+LOGGER = logging.getLogger(__name__)
+
+
 class EndpointHandler():
     def __init__(self, path=""):
         config = PeftConfig.from_pretrained(path)
@@ -14,35 +18,21 @@ class EndpointHandler():
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
-       data args:
-            prompt (:obj:`str`):
-            temperature (:obj:`float`, `optional`, defaults to 0.5):
-            eos_token_id (:obj:`int`, `optional`, defaults to tokenizer.eos_token_id):
-            early_stopping (:obj:`bool`, `optional`, defaults to `True`):
-            repetition_penalty (:obj:`float`, `optional`, defaults to 0.3):
-      Return:
-            A :obj:`str` : generated sequences
+        Args:
+            data (Dict): The payload with the text prompt and generation parameters.
         """
+        LOGGER.info(f"Received data: {data}")
         # Get inputs
-        prompt = data.pop("prompt", None)
-        temperature = data.pop("temperature", 0.5)
-        eos_token_id = data.pop("eos_token_id", self.tokenizer.eos_token_id)
-        early_stopping = data.pop('early_stopping', True)
-        repetition_penalty = data.pop('repetition_penalty', 0.3)
-        max_new_tokens = data.pop('max_new_tokens', 100)
-
-        if prompt is None:
-            raise ValueError("No prompt provided.")
-        
-        # Run prediction
+        prompt = data.pop("prompt", data)
+        parameters = data.pop("parameters", None)
+        # Preprocess
         inputs = self.tokenizer(prompt, return_tensors="pt")
-        prediction = self.model.generate(
-            **inputs,
-            temperature=temperature,
-            eos_token_id=eos_token_id,
-            early_stopping=early_stopping,
-            repetition_penalty=repetition_penalty,
-            max_new_tokens=max_new_tokens
-        )
-
-        return prediction
+        # Forward
+        if parameters is not None:
+            outputs = self.model.generate(**inputs, **parameters)
+        else:
+            outputs = self.model.generate(**inputs)
+        # Postprocess
+        prediction = self.tokenizer.decode(outputs[0])
+        LOGGER.info(f"Generated text: {prediction}")
+        return [{"generated_text": prediction}]