add custom handler and modify pipeline

handler.py

@@ -7,6 +7,7 @@ from transformers.pipelines import PIPELINE_REGISTRY
 from bibert_multitask_classification import BiBert_MultiTaskPipeline
 from bert_for_sequence_classification import BertForSequenceClassification
 from transformers.utils import logging
+from time import perf_counter
 
 
 PIPELINE_REGISTRY.register_pipeline("bibert-multitask-classification", pipeline_class=BiBert_MultiTaskPipeline, pt_model=BertForSequenceClassification)
@@ -32,12 +33,17 @@ tasks = [
 class EndpointHandler():
   def __init__(self, path=""):
     # Preload all the elements you are going to need at inference.
-    tokenizer = AutoTokenizer.from_pretrained(path)
+    logger.info("The device is %s.", device)
+
+    t0 = perf_counter()
 
+    tokenizer = AutoTokenizer.from_pretrained(path)
     model = BertForSequenceClassification.from_pretrained(path, tasks_map=tasks).to(device)
-    
     self.classifier_s = pipeline("bibert-multitask-classification", model = model, task_id="0", tokenizer=tokenizer, device = device)
     self.classifier_p = pipeline("bibert-multitask-classification", model = model, task_id="1", tokenizer=tokenizer, device = device)
+    elapsed = 1000 * (perf_counter() - t0)
+    logger.info("Models and tokenizer Polarity loaded in %d ms.", elapsed)
+
 
   def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
     """
@@ -50,20 +56,21 @@ class EndpointHandler():
     
     inputs = data.pop("inputs", data)
     lang = data.pop("lang", None)
-    logger.info(inputs)
-    logger.info(lang)
+    logger.info("The language of Verbatim is %s.", lang)
     if isinstance(inputs, str):
       inputs = [inputs]
     
+    t0 = perf_counter()
+    prediction_res = []
     prediction_p = self.classifier_p(inputs) 
     label = prediction_p[0]['label']
     score = prediction_p[0]['score']
-    
+
     if label == '0' and score >= 0.75:
       label = 2
-      return [{"label":label, "score": score}]
+      prediction_res = [{"label":label, "score": score}]
     else:
-      prediction_s = self.classifier_s(inputs)
-      label = prediction_s[0]['label']
-      score = prediction_s[0]['score']
-    return prediction_s
+      prediction_res = self.classifier_s(inputs)
+    elapsed = 1000 * (perf_counter() - t0)
+    logger.info("Model prediction time: %d ms.", elapsed)  
+    return prediction_res