Upload folder using huggingface_hub

handler.py

@@ -2,7 +2,6 @@ from typing import Dict, List, Any
 import torch
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import logging
-logger = logging.getLogger(__name__)
 
 
 class EndpointHandler:
@@ -11,9 +10,14 @@ class EndpointHandler:
         self.model =  AutoModelForCausalLM.from_pretrained(path, device_map="auto", load_in_8bit=True)
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.pipe = pipeline("text-generation", self.model, self.tokenizer)
+        logging.basicConfig(
+            level=logging.INFO,
+            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+        )
+        self.logger = logging.getLogger(__name__)
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
-        logger.info("CALL DATA:", data)
+        self.logger.info("CALL DATA:", data)
 
         # process input
         inputs = data.pop("inputs", data)
@@ -25,6 +29,6 @@ class EndpointHandler:
             output = self.pipe(inputs)
 
         response = output[0]['generated_text'][-1]
-        logger.info("RESPONSE:", response)
+        self.logger.info("RESPONSE:", response)
 
         return {"generated_text": response}