zjiayao commited on
Commit
c316f5e
1 Parent(s): f159c9b

Create temp_predictor.py

Browse files
Files changed (1) hide show
  1. src/temp_predictor.py +76 -0
src/temp_predictor.py ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import spacy
2
+ import transformers
3
+ import numpy as np
4
+
5
+ class TempPredictor:
6
+ def __init__(self, model, tokenizer, device,
7
+ spacy_model="en_core_web_sm"):
8
+ self._model = model
9
+ self._model.to(device)
10
+ self._model.eval()
11
+ self._tokenizer = tokenizer
12
+ self._mtoken = self._tokenizer.mask_token
13
+ self.unmasker = transformers.pipeline("fill-mask", model=self._model, tokenizer=self._tokenizer, device=0)
14
+ try:
15
+ self._spacy = spacy.load(spacy_model)
16
+ except Exception as e:
17
+ self._spacy = spacy.load("en_core_web_sm")
18
+ print(f"Failed to load spacy model {spacy_model}, use default 'en_core_web_sm'\n{e}")
19
+
20
+ def _extract_token_prob(self, arr, token, crop=1):
21
+ for it in arr:
22
+ if len(it["token_str"]) >= crop and (token == it["token_str"][crop:]):
23
+ return it["score"]
24
+ return 0.
25
+
26
+ def _sent_lowercase(self, s):
27
+ try:
28
+ return s[0].lower() + s[1:]
29
+ except:
30
+ return s
31
+
32
+ def _remove_punct(self, s):
33
+ try:
34
+ return s[:-1]
35
+ except:
36
+ return s
37
+
38
+ def predict(self, e1, e2, top_k=5):
39
+ txt = self._remove_punct(e1) + " " + self._mtoken + " " + self._sent_lowercase(e2)
40
+ return self.unmasker(txt, top_k=top_k)
41
+
42
+ def batch_predict(self, instances, top_k=5):
43
+ txt = [self._remove_punct(e1) + " " + self._mtoken + " " + self._sent_lowercase(e2)
44
+ for (e1, e2) in instances]
45
+ return self.unmasker(txt, top_k=top_k)
46
+
47
+
48
+ def get_temp(self, e1, e2, top_k=5, crop=1):
49
+ inst1 = self.predict(e1, e2, top_k)
50
+ inst2 = self.predict(e2, e1, top_k)
51
+
52
+ # e1 before e2
53
+ b1 = self._extract_token_prob(inst1, "before", crop=crop)
54
+ b2 = self._extract_token_prob(inst2, "after", crop=crop)
55
+
56
+ # e1 after e2
57
+ a1 = self._extract_token_prob(inst1, "after", crop=crop)
58
+ a2 = self._extract_token_prob(inst2, "before", crop=crop)
59
+
60
+ return (b1+b2)/2, (a1+a2)/2
61
+
62
+ def get_temp_batch(self, instances, top_k=5, crop=1):
63
+ reverse_instances = [(e2, e1) for (e1, e2) in instances]
64
+ fwd_preds = self.batch_predict(instances, top_k=top_k)
65
+ bwd_preds = self.batch_predict(reverse_instances, top_k=top_k)
66
+
67
+ b1s = np.array([ self._extract_token_prob(pred, "before", crop=crop) for pred in fwd_preds ])
68
+ b2s = np.array([ self._extract_token_prob(pred, "before", crop=crop) for pred in bwd_preds ])
69
+ a1s = np.array([ self._extract_token_prob(pred, "after", crop=crop) for pred in fwd_preds ])
70
+ a2s = np.array([ self._extract_token_prob(pred, "after", crop=crop) for pred in bwd_preds ])
71
+
72
+ return np.array([np.array(b1s+b2s)/2, np.array(a1s+a2s)/2]).T
73
+
74
+
75
+ def __call__(self, *args, **kwargs):
76
+ return self.get_temp(*args, **kwargs)