Update handler.py

handler.py

@@ -1,4 +1,3 @@
-from dataclasses import dataclass
 import torch
 import numpy as np
 import transformers
@@ -208,29 +207,6 @@ class BiBert_MultiTaskPipeline(Pipeline):
           dict_scores = dict_scores[:top_k]
     return dict_scores
 
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-checkpoint="HCKLab/BiBert-MultiTask-1"
-
-tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-
-tasks = [
-    Task(id=0, name='label_classification', type='seq_classification', num_labels=5),
-    Task(id=1, name='binary_classification', type='seq_classification', num_labels=2)
-    ]
-
-multitask_model = BertForSequenceClassification.from_pretrained(
-        checkpoint,
-        tasks_map=tasks
-    ).to(device)
-
-
-PIPELINE_REGISTRY.register_pipeline(
-    "bibert-multitask-classification",
-    pipeline_class=BiBert_MultiTaskPipeline,
-    pt_model=BertForSequenceClassification
-)
-
 class EndpointHandler():
   def __init__(self, path=""):
     # Preload all the elements you are going to need at inference.
@@ -239,6 +215,19 @@ class EndpointHandler():
     pipeline_class=BiBert_MultiTaskPipeline,
     pt_model=BertForSequenceClassification
     )
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+
+    tasks = [
+        Task(id=0, name='label_classification', type='seq_classification', num_labels=5),
+        Task(id=1, name='binary_classification', type='seq_classification', num_labels=2)
+        ]
+
+    multitask_model = BertForSequenceClassification.from_pretrained(
+            path,
+            tasks_map=tasks
+        ).to(device)
     self.classifier_s = pipeline("bibert-multitask-classification", model = multitask_model, task_id="0", tokenizer=tokenizer, device = device)
     self.classifier_p = pipeline("bibert-multitask-classification", model = multitask_model, task_id="1", tokenizer=tokenizer, device = device)