changes for testin

handler.py

@@ -1,30 +1,30 @@
 from typing import Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
+device = "cuda"
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # Load the model
+        # load the model
         tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B-Instruct")
         model = AutoModelForCausalLM.from_pretrained(
             "Qwen/Qwen2-1.5B-Instruct",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="cuda" if torch.cuda.is_available() else "auto" # Include device_map for correct device allocation
+            torch_dtype="auto",
+            device_map="auto"
         )
-
-        # Create inference pipeline without specifying the device
+        # create inference pipeline
         self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-    def __call__(self, data: Any) -> List[List[Dict[str, Any]]]:
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         inputs = data.pop("inputs", data)
-        parameters = data.pop("parameters", {})
-
-        if isinstance(inputs, str):
-            inputs = [inputs]
+        parameters = data.pop("parameters", None)
 
-        # Get predictions from the pipeline
-        prediction = self.pipeline(inputs, **parameters)
+        # 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
 
 # Example usage