Update handler.py

handler.py

@@ -1,19 +1,28 @@
 from typing import Dict, Any, List
 import torch
+import logging
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 class EndpointHandler():
     def __init__(self, path=""):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         try:
+            logger.info(f"Loading model and tokenizer from path: {path}")
             self.model = AutoModelForSeq2SeqLM.from_pretrained(f"{path}/model_v2/").to(self.device)
             self.tokenizer = AutoTokenizer.from_pretrained(f"{path}/model_v2/")
         except Exception as e:
-            print(f"Error loading model or tokenizer from path {path}: {e}")
+            logger.error(f"Error loading model or tokenizer from path {path}: {e}")
             # Handle error (e.g., exit or set model/tokenizer to None)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        inputs = data.get("inputs", "")
+        if self.model is None or self.tokenizer is None:
+            error_message = "Model or tokenizer not properly initialized"
+            logger.error(error_message)
+            return [{"error": error_message}]
+        inputs = data.get("inputs")
         if not inputs:
             return [{"error": "No inputs provided"}]