forked from weaviate/t2v-transformers-models
-
Notifications
You must be signed in to change notification settings - Fork 0
/
vectorizer.py
231 lines (181 loc) · 8.45 KB
/
vectorizer.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
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
import asyncio
from concurrent.futures import ThreadPoolExecutor
import math
from typing import Optional
import torch
import nltk
from nltk.tokenize import sent_tokenize
from pydantic import BaseModel
from transformers import (
AutoModel,
AutoTokenizer,
T5ForConditionalGeneration,
T5Tokenizer,
DPRContextEncoder,
DPRQuestionEncoder,
)
# limit transformer batch size to limit parallel inference, otherwise we run
# into memory problems
MAX_BATCH_SIZE = 25 # TODO: take from config
DEFAULT_POOL_METHOD="masked_mean"
class VectorInputConfig(BaseModel):
pooling_strategy: str
class VectorInput(BaseModel):
text: str
config: Optional[VectorInputConfig] = None
class Vectorizer:
model: AutoModel
tokenizer: AutoTokenizer
cuda: bool
cuda_core: str
model_type: str
direct_tokenize: bool
executor: ThreadPoolExecutor
def __init__(self, model_path: str, cuda_support: bool, cuda_core: str, cuda_per_process_memory_fraction: float, model_type: str, architecture: str, direct_tokenize: bool):
self.cuda = cuda_support
self.cuda_core = cuda_core
self.cuda_per_process_memory_fraction = cuda_per_process_memory_fraction
self.model_type = model_type
self.direct_tokenize = direct_tokenize
self.model_delegate: HFModel = ModelFactory.model(model_type, architecture, cuda_support, cuda_core)
self.model = self.model_delegate.create_model(model_path)
if self.cuda:
self.model.to(self.cuda_core)
if self.cuda_per_process_memory_fraction:
torch.cuda.set_per_process_memory_fraction(self.cuda_per_process_memory_fraction)
self.model.eval() # make sure we're in inference mode, not training
self.tokenizer = self.model_delegate.create_tokenizer(model_path)
self.executor = ThreadPoolExecutor()
nltk.data.path.append('./nltk_data')
def tokenize(self, text:str):
return self.tokenizer(text, padding=True, truncation=True, max_length=500,
add_special_tokens = True, return_tensors="pt")
def get_embeddings(self, batch_results):
return self.model_delegate.get_embeddings(batch_results)
def get_batch_results(self, tokens, text):
return self.model_delegate.get_batch_results(tokens, text)
def pool_embedding(self, batch_results, tokens, config):
return self.model_delegate.pool_embedding(batch_results, tokens, config)
def _vectorize(self, text: str, config: VectorInputConfig):
with torch.no_grad():
if self.direct_tokenize:
# create embeddings without tokenizing text
tokens = self.tokenize(text)
if self.cuda:
tokens.to(self.cuda_core)
batch_results = self.get_batch_results(tokens, text)
batch_sum_vectors = self.pool_embedding(batch_results, tokens, config)
return batch_sum_vectors.detach()
else:
# tokenize text
sentences = sent_tokenize(' '.join(text.split(),))
num_sentences = len(sentences)
number_of_batch_vectors = math.ceil(num_sentences / MAX_BATCH_SIZE)
batch_sum_vectors = 0
for i in range(0, number_of_batch_vectors):
start_index = i * MAX_BATCH_SIZE
end_index = start_index + MAX_BATCH_SIZE
tokens = self.tokenize(sentences[start_index:end_index])
if self.cuda:
tokens.to(self.cuda_core)
batch_results = self.get_batch_results(tokens, sentences[start_index:end_index])
batch_sum_vectors += self.pool_embedding(batch_results, tokens, config)
return batch_sum_vectors.detach() / num_sentences
async def vectorize(self, text: str, config: VectorInputConfig):
return await asyncio.wrap_future(self.executor.submit(self._vectorize, text, config))
class HFModel:
def __init__(self, cuda_support: bool, cuda_core: str):
super().__init__()
self.model = None
self.tokenizer = None
self.cuda = cuda_support
self.cuda_core = cuda_core
def create_tokenizer(self, model_path):
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
return self.tokenizer
def create_model(self, model_path):
self.model = AutoModel.from_pretrained(model_path)
return self.model
def get_embeddings(self, batch_results):
return batch_results[0]
def get_batch_results(self, tokens, text):
return self.model(**tokens)
def pool_embedding(self, batch_results, tokens, config: VectorInputConfig):
pooling_method = self.pool_method_from_config(config)
if pooling_method == "cls":
return self.get_embeddings(batch_results)[:, 0, :].sum(0)
elif pooling_method == "masked_mean":
return self.pool_sum(self.get_embeddings(batch_results), tokens['attention_mask'])
else:
raise Exception(f"invalid pooling method '{pooling_method}'")
def pool_method_from_config(self, config: VectorInputConfig):
if config is None:
return DEFAULT_POOL_METHOD
if config.pooling_strategy is None or config.pooling_strategy == "":
return DEFAULT_POOL_METHOD
return config.pooling_strategy
def get_sum_embeddings_mask(self, embeddings, input_mask_expanded):
if self.cuda:
sum_embeddings = torch.sum(embeddings * input_mask_expanded, 1).to(self.cuda_core)
sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9).to(self.cuda_core)
return sum_embeddings, sum_mask
else:
sum_embeddings = torch.sum(embeddings * input_mask_expanded, 1)
sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
return sum_embeddings, sum_mask
def pool_sum(self, embeddings, attention_mask):
input_mask_expanded = attention_mask.unsqueeze(-1).expand(embeddings.size()).float()
sum_embeddings, sum_mask = self.get_sum_embeddings_mask(embeddings, input_mask_expanded)
sentences = sum_embeddings / sum_mask
return sentences.sum(0)
class DPRModel(HFModel):
def __init__(self, architecture: str, cuda_support: bool, cuda_core: str):
super().__init__(cuda_support, cuda_core)
self.model = None
self.architecture = architecture
def create_model(self, model_path):
if self.architecture == "DPRQuestionEncoder":
self.model = DPRQuestionEncoder.from_pretrained(model_path)
else:
self.model = DPRContextEncoder.from_pretrained(model_path)
return self.model
def get_batch_results(self, tokens, text):
return self.model(tokens['input_ids'], tokens['attention_mask'])
def pool_embedding(self, batch_results, tokens, config: VectorInputConfig):
# no pooling needed for DPR
return batch_results["pooler_output"][0]
class T5Model(HFModel):
def __init__(self, cuda_support: bool, cuda_core: str):
super().__init__(cuda_support, cuda_core)
self.model = None
self.tokenizer = None
self.cuda = cuda_support
self.cuda_core = cuda_core
def create_model(self, model_path):
self.model = T5ForConditionalGeneration.from_pretrained(model_path)
return self.model
def create_tokenizer(self, model_path):
self.tokenizer = T5Tokenizer.from_pretrained(model_path)
return self.tokenizer
def get_embeddings(self, batch_results):
return batch_results["encoder_last_hidden_state"]
def get_batch_results(self, tokens, text):
input_ids, attention_mask = tokens['input_ids'], tokens['attention_mask']
target_encoding = self.tokenizer(
text, padding="longest", max_length=500, truncation=True
)
labels = target_encoding.input_ids
if self.cuda:
labels = torch.tensor(labels).to(self.cuda_core)
else:
labels = torch.tensor(labels)
return self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
class ModelFactory:
@staticmethod
def model(model_type, architecture, cuda_support: bool, cuda_core: str):
if model_type == 't5':
return T5Model(cuda_support, cuda_core)
elif model_type == 'dpr':
return DPRModel(architecture, cuda_support, cuda_core)
else:
return HFModel(cuda_support, cuda_core)