update handler for inputs and parameters

handler.py

@@ -2,21 +2,26 @@ from typing import Dict, List, Any
 import transformers
 import torch
 
-MAX_TOKENS=8192
+MAX_TOKENS=1024
 
 class EndpointHandler(object):
     def __init__(self, path=''):
         self.pipeline: transformers.Pipeline = transformers.pipeline(
             "text-generation",
             model="humane-intelligence/gemma2-9b-cpt-sealionv3-instruct-endpoint",
-            model_kwargs={"torch_dtype": torch.bfloat16, "low_cpu_mem_usage": True, },
+            model_kwargs={"torch_dtype": torch.bfloat16 },
             device_map="auto",
         )
 
-    def __call__(self, text_inputs: Any) -> List[List[Dict[str, float]]]:
-        outputs = self.pipeline(
-            text_inputs,
-            max_new_tokens=MAX_TOKENS,
-        )
-        print(outputs[0]["generated_text"][-1])
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        inputs = data.pop("inputs")
+
+        if parameters:= data.pop("parameters", None):
+            outputs = self.pipeline(
+                inputs,
+                **parameters
+            )
+        else:
+            outputs = self.pipeline(inputs, max_new_tokens=MAX_TOKENS)
+
         return outputs