/
generation.py
193 lines (163 loc) · 6.43 KB
/
generation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
from typing import Dict, List, Optional, Tuple
import torch
from fire import Fire
from torch import Tensor
from transformers import PreTrainedModel, PreTrainedTokenizerFast
from encoding import ExtractEncoder
from utils import DynamicModel, RelationSentence, find_sublist_index
class TextGenerator(DynamicModel):
model: PreTrainedModel
tokenizer: PreTrainedTokenizerFast
scores: Optional[List[Tensor]] = None
max_length: int
def tokenize(self, texts: List[str], **kwargs):
return self.tokenizer(
texts,
padding=True,
truncation=True,
max_length=self.max_length,
return_tensors="pt",
**kwargs,
).to(self.model.device)
def run(
self,
texts: List[str],
do_sample=True,
top_k=50,
temperature=1.0,
num_return: int = 4,
prompt: Optional[str] = None,
prompt_ids: Optional[List[int]] = None,
multi_prompt_ids: Optional[List[List[int]]] = None,
decoder_input_ids: Optional[Tensor] = None,
save_scores: bool = False,
**kwargs,
) -> List[str]:
# https://huggingface.co/transformers/v4.7.0/main_classes/model.html#generation
tok = self.tokenizer
eos, bos = tok.eos_token_id, tok.bos_token_id
if prompt is not None:
prompt_ids = self.tokenizer(prompt, add_special_tokens=False).input_ids
if prompt_ids is not None:
prompt_ids = [eos, bos] + prompt_ids
decoder_input_ids = torch.tensor([prompt_ids])
if multi_prompt_ids is not None:
assert len(texts) == len(multi_prompt_ids)
multi_prompt_ids = [[eos, bos] + lst for lst in multi_prompt_ids]
decoder_input_ids = torch.tensor(multi_prompt_ids)
if decoder_input_ids is not None:
kwargs.update(decoder_input_ids=decoder_input_ids.to(self.model.device))
outputs = self.model.generate(
**self.tokenize(texts),
do_sample=do_sample,
top_k=top_k,
temperature=temperature,
num_return_sequences=num_return,
return_dict_in_generate=True,
output_scores=save_scores,
max_length=self.max_length,
**kwargs,
)
self.scores = None
if save_scores:
self.scores = [_ for _ in torch.stack(outputs.scores, 1).cpu()]
return self.decode(outputs.sequences)
def decode(self, outputs) -> List[str]:
tok = self.tokenizer
texts = tok.batch_decode(
outputs, skip_special_tokens=False, clean_up_tokenization_spaces=False
)
# Manually remove <bos><eos><pad> in case we have custom special tokens
special_tokens = [tok.eos_token, tok.bos_token, tok.pad_token]
for i, t in enumerate(texts):
for token in special_tokens:
t = t.replace(token, "")
texts[i] = t
return texts
class LabelConstraint:
def __init__(
self,
labels: List[str],
tokenizer: PreTrainedTokenizerFast,
prefix: str = " Relation :",
):
self.prefix: List[int] = tokenizer(prefix, add_special_tokens=False).input_ids
self.label_map: Dict[int, str] = {
tokenizer(" " + x, add_special_tokens=False).input_ids[0]: x for x in labels
}
self.tokenizer = tokenizer
def run(self, triplet: RelationSentence, scores: Tensor) -> RelationSentence:
triplet = triplet.copy(deep=True)
assert scores.ndim == 2
token_ids = scores.argmax(dim=-1).int().tolist()
i = find_sublist_index(token_ids, self.prefix)
if i == -1:
return triplet
position = i + len(self.prefix)
best = ""
best_score = -1e9
for j, label in self.label_map.items():
score = scores[position, j].item()
if score > best_score:
best = label
best_score = score
if triplet.label in self.label_map.values():
assert best == triplet.label
assert len(best) > 0
triplet.label = best
triplet.score = best_score
return triplet
class TripletSearchDecoder(DynamicModel):
gen: TextGenerator
constraint: LabelConstraint
encoder: ExtractEncoder
top_k: int = 4
def generate(self, text: str, **kwargs) -> Tuple[str, Tensor]:
outputs = self.gen.run(
[text],
do_sample=False,
num_return=1,
num_beams=1,
save_scores=True,
**kwargs,
)
assert len(outputs) == 1
assert self.gen.scores is not None
scores = torch.log_softmax(self.gen.scores[0], dim=-1)
assert scores.ndim == 2
return outputs[0], scores
def find_prefix_end(self, token_ids: List[str], prefix: str) -> int:
prefix_ids = self.gen.tokenizer(prefix, add_special_tokens=False).input_ids
i = find_sublist_index(token_ids, prefix_ids)
position = i + len(prefix_ids)
return position
def branch(
self, text: str, prefix: str, prompt: Optional[str] = None, **kwargs
) -> List[Tuple[str, float]]:
_, scores = self.generate(text, prompt=prompt, **kwargs)
token_ids = scores.argmax(dim=-1).int().tolist()
i = self.find_prefix_end(token_ids, prefix)
pairs = []
for j in torch.argsort(scores[i])[-self.top_k :]:
p = (prompt or "") + self.gen.decode([token_ids[:i] + [j]])[0]
pairs.append((p, scores[i, j].item()))
return pairs
def run(self, text: str) -> List[RelationSentence]:
x = self.encoder.encode_x(text)
outputs = []
for prompt_a, score_a in self.branch(x, prefix="Head Entity :"):
for prompt_b, score_b in self.branch(
x, prefix=" Tail Entity :", prompt=prompt_a
):
output, scores = self.generate(x, prompt=prompt_b)
token_ids = token_ids = scores.argmax(dim=-1).int().tolist()
i = self.find_prefix_end(token_ids, prefix=" Relation :")
score_c = max(scores[i].tolist())
s = self.encoder.safe_decode(x=x, y=output)
s = self.constraint.run(s, scores)
# score_c = s.score # From LabelConstraint
s.score = (score_a + score_b + score_c) / 3
outputs.append(s)
return outputs
if __name__ == "__main__":
Fire()