add logging to handler

handler.py

@@ -1,6 +1,10 @@
 from typing import Dict, List, Any
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import logging
+
+logging.basicConfig(level=logging.INFO)
+LOGGER = logging.getLogger(__name__)
 
 
 # Prompts for the different tasks
@@ -29,6 +33,10 @@ class EndpointHandler:
         self.model = AutoModelForCausalLM.from_pretrained(
             path, device_map="auto", torch_dtype=torch.bfloat16
         )
+        LOGGER.info(f"Inference model loaded from {path}")
+        LOGGER.info(f"Model outline: {self.model}")
+        LOGGER.info(f"Model device: {self.model.device}")
+
         # Fix the pad and bos tokens to avoid bug in the tokenizer
         pad_token = "<unk>"
         bos_token = "<|endoftext|>"
@@ -90,6 +98,7 @@ class EndpointHandler:
         Return:
               A :obj:`list` | `dict`: will be serialized and returned
         """
+        LOGGER.info(f"Received data: {data}")
 
         # Get inputs
         input_a = data.pop("input_a", None)