Update handler.py

handler.py

@@ -3,12 +3,17 @@ import transformers
 import torch
 from datetime import datetime
 from transformers import StoppingCriteria, StoppingCriteriaList
+from transformers.utils import logging
+
+logging.set_verbosity_info()
+logger = logging.get_logger("transformers")
 
 
 class EndpointHandler():
 
     def __init__(self, path=""):
-        print(f"Hugging face handler path {path}")
+        logger.info("111111111111111111111111111")
+        logger.info(f"Hugging face handler path {path}")
         path = 'mosaicml/mpt-7b'
         self.model = transformers.AutoModelForCausalLM.from_pretrained(path,
             #"/Users/itamarlevi/Downloads/my_repo_hf/hf/mpt-7b/venv/Itamarl/test",
@@ -49,5 +54,6 @@ class EndpointHandler():
 
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:    
-        res = self.generate_text("Explain to me the difference between nuclear fission and fusion.")
+        logger.info(f"iiinnnnnnnnnn    {data}")
+        res = self.generate_text(data)
         return res