update handler

handler.py

@@ -4,34 +4,37 @@ import torch
 
 class EndpointHandler:
     def __init__(self, path=""):
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = 0 if torch.cuda.is_available() else -1  # 0 for GPU, -1 for CPU
         
-        # 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 device == "cuda" else torch.float32,
-            device_map="auto"
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            device_map="cuda" # for single instance one GPU
         )
         
-        # create inference pipeline without specifying the device
-        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        # Create inference pipeline with the correct device
+        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
 
-    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+    def __call__(self, data: Any) -> List[List[Dict[str, Any]]]:
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
 
-        # Ensure inputs are on the GPU if available
         if isinstance(inputs, str):
             inputs = [inputs]
 
-        # Tensor input handling
-        try:
-            inputs = torch.tensor(inputs).cuda() if torch.cuda.is_available() else torch.tensor(inputs)
-        except:
-            pass  # If inputs are not tensors (e.g., strings), continue without conversion
-
-        # pass inputs with all kwargs in data
+        # Get predictions from the pipeline
         prediction = self.pipeline(inputs, **parameters)
         
         return prediction
+
+# Example usage
+if __name__ == "__main__":
+    handler = EndpointHandler()
+    data = {
+        "inputs": "Hello, how can I",
+        "parameters": {"max_length": 50, "num_return_sequences": 1}
+    }
+    result = handler(data)
+    print(result)