davanstrien HF staff commited on
Commit
9dcab23
1 Parent(s): b903e65

Create pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +48 -0
pipeline.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ import os
3
+
4
+ import json
5
+ import numpy as np
6
+
7
+ from fastai.learner import load_learner
8
+
9
+ class PreTrainedPipeline():
10
+
11
+ def __init__(self, path=""):
12
+
13
+ # IMPLEMENT_THIS
14
+
15
+ # Preload all the elements you are going to need at inference.
16
+
17
+ # For instance your model, processors, tokenizer that might be needed.
18
+
19
+ # This function is only called once, so do all the heavy processing I/O here"""
20
+
21
+ self.model = load_learner(os.path.join(path, "20211115-model.pkl"))
22
+
23
+ with open(os.path.join(path, "config.json")) as config:
24
+
25
+ config = json.load(config)
26
+
27
+ self.id2label = config["id2label"]
28
+
29
+ def __call__(self, inputs: str) -> List[Dict[str, Any]]:
30
+
31
+
32
+ # IMPLEMENT_THIS
33
+
34
+
35
+
36
+ _, _, preds = self.model.predict(inputs)
37
+
38
+ preds = preds.tolist()
39
+
40
+ labels = [
41
+
42
+ {"label": str(self.id2label["0"]), "score": preds[0]},
43
+
44
+ {"label": str(self.id2label["1"]), "score": preds[1]},
45
+
46
+ ]
47
+
48
+ return labels