Kalaoke commited on
Commit
5092549
·
1 Parent(s): 943a701

Delete handler.py

Browse files
Files changed (1) hide show
  1. handler.py +0 -259
handler.py DELETED
@@ -1,259 +0,0 @@
1
- import torch
2
- import numpy as np
3
- import transformers
4
- from torch import nn
5
- from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
6
- from typing import List, Optional, Tuple, Union, Dict, Any
7
-
8
- from transformers import BertTokenizer, Pipeline
9
- from transformers import models, DataCollatorWithPadding, AutoTokenizer
10
- from transformers.modeling_outputs import SequenceClassifierOutput
11
- from transformers.pipelines import PIPELINE_REGISTRY
12
-
13
- from transformers.models.bert.configuration_bert import BertConfig
14
- from transformers.models.bert.modeling_bert import (
15
- BertPreTrainedModel,
16
- BERT_INPUTS_DOCSTRING,
17
- _TOKENIZER_FOR_DOC,
18
- _CHECKPOINT_FOR_DOC,
19
- BERT_START_DOCSTRING,
20
- _CONFIG_FOR_DOC,
21
- _SEQ_CLASS_EXPECTED_OUTPUT,
22
- _SEQ_CLASS_EXPECTED_LOSS,
23
- BertModel,
24
- )
25
-
26
- from transformers.file_utils import (
27
- add_code_sample_docstrings,
28
- add_start_docstrings_to_model_forward,
29
- add_start_docstrings
30
- )
31
-
32
- @dataclass
33
- class Task:
34
- id: int
35
- name: str
36
- type: str
37
- num_labels: int
38
-
39
-
40
- @add_start_docstrings(
41
- """
42
- Bert Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
43
- output) e.g. for GLUE tasks.
44
- """,
45
- BERT_START_DOCSTRING,
46
- )
47
- class BertForSequenceClassification(BertPreTrainedModel):
48
- def __init__(self, config, **kwargs):
49
- super().__init__(transformers.PretrainedConfig())
50
- #task_labels_map={"binary_classification": 2, "label_classification": 5}
51
- self.tasks = kwargs.get("tasks_map", {})
52
- self.config = config
53
-
54
- self.bert = BertModel(config)
55
- classifier_dropout = (
56
- config.classifier_dropout
57
- if config.classifier_dropout is not None
58
- else config.hidden_dropout_prob
59
- )
60
- self.dropout = nn.Dropout(classifier_dropout)
61
- ## add task specific output heads
62
- self.classifier1 = nn.Linear(
63
- config.hidden_size, self.tasks[0].num_labels
64
- )
65
- self.classifier2 = nn.Linear(
66
- config.hidden_size, self.tasks[1].num_labels
67
- )
68
-
69
- self.init_weights()
70
-
71
- @add_start_docstrings_to_model_forward(
72
- BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length")
73
- )
74
- @add_code_sample_docstrings(
75
- processor_class=_TOKENIZER_FOR_DOC,
76
- checkpoint=_CHECKPOINT_FOR_DOC,
77
- output_type=SequenceClassifierOutput,
78
- config_class=_CONFIG_FOR_DOC,
79
- expected_output=_SEQ_CLASS_EXPECTED_OUTPUT,
80
- expected_loss=_SEQ_CLASS_EXPECTED_LOSS,
81
- )
82
- def forward(
83
- self,
84
- input_ids: Optional[torch.Tensor] = None,
85
- attention_mask: Optional[torch.Tensor] = None,
86
- token_type_ids: Optional[torch.Tensor] = None,
87
- position_ids: Optional[torch.Tensor] = None,
88
- head_mask: Optional[torch.Tensor] = None,
89
- inputs_embeds: Optional[torch.Tensor] = None,
90
- labels: Optional[torch.Tensor] = None,
91
- output_attentions: Optional[bool] = None,
92
- output_hidden_states: Optional[bool] = None,
93
- return_dict: Optional[bool] = None,
94
- task_ids=None,
95
- ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
96
- r"""
97
- labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
98
- Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
99
- config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
100
- If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
101
- """
102
- return_dict = (
103
- return_dict if return_dict is not None else self.config.use_return_dict
104
- )
105
-
106
- outputs = self.bert(
107
- input_ids,
108
- attention_mask=attention_mask,
109
- token_type_ids=token_type_ids,
110
- position_ids=position_ids,
111
- head_mask=head_mask,
112
- inputs_embeds=inputs_embeds,
113
- output_attentions=output_attentions,
114
- output_hidden_states=output_hidden_states,
115
- return_dict=return_dict,
116
- )
117
-
118
- pooled_output = outputs[1]
119
-
120
- pooled_output = self.dropout(pooled_output)
121
-
122
- unique_task_ids_list = torch.unique(task_ids).tolist()
123
- loss_list = []
124
- logits = None
125
- for unique_task_id in unique_task_ids_list:
126
-
127
- loss = None
128
- task_id_filter = task_ids == unique_task_id
129
-
130
- if unique_task_id == 0:
131
- logits = self.classifier1(pooled_output[task_id_filter])
132
- elif unique_task_id == 1:
133
- logits = self.classifier2(pooled_output[task_id_filter])
134
-
135
-
136
- if labels is not None:
137
- loss_fct = CrossEntropyLoss()
138
- loss = loss_fct(logits.view(-1, self.tasks[unique_task_id].num_labels), labels[task_id_filter].view(-1))
139
- loss_list.append(loss)
140
-
141
- # logits are only used for eval. and in case of eval the batch is not multi task
142
- # For training only the loss is used
143
-
144
- if loss_list:
145
- loss = torch.stack(loss_list).mean()
146
- if not return_dict:
147
- output = (logits,) + outputs[2:]
148
- return ((loss,) + output) if loss is not None else output
149
-
150
- return SequenceClassifierOutput(
151
- loss=loss,
152
- logits=logits,
153
- hidden_states=outputs.hidden_states,
154
- attentions=outputs.attentions,
155
- )
156
-
157
- def softmax(_outputs):
158
- maxes = np.max(_outputs, axis=-1, keepdims=True)
159
- shifted_exp = np.exp(_outputs - maxes)
160
- return shifted_exp / shifted_exp.sum(axis=-1, keepdims=True)
161
-
162
- class BiBert_MultiTaskPipeline(Pipeline):
163
-
164
-
165
- def _sanitize_parameters(self, **kwargs):
166
-
167
- preprocess_kwargs = {}
168
- if "task_id" in kwargs:
169
- preprocess_kwargs["task_id"] = kwargs["task_id"]
170
-
171
- forward_kwargs = {}
172
- if "task_id" in kwargs:
173
- forward_kwargs["task_id"] = kwargs["task_id"]
174
-
175
- postprocess_kwargs = {}
176
- if "top_k" in kwargs:
177
- postprocess_kwargs["top_k"] = kwargs["top_k"]
178
- postprocess_kwargs["_legacy"] = False
179
- return preprocess_kwargs, forward_kwargs, postprocess_kwargs
180
-
181
-
182
-
183
- def preprocess(self, inputs, task_id):
184
- return_tensors = self.framework
185
- feature = self.tokenizer(inputs, padding = True, return_tensors=return_tensors).to(self.device)
186
- task_ids = np.full(shape=1,fill_value=task_id, dtype=int)
187
- feature["task_ids"] = torch.IntTensor(task_ids)
188
- return feature
189
-
190
- def _forward(self, model_inputs, task_id):
191
- return self.model(**model_inputs)
192
-
193
- def postprocess(self, model_outputs, top_k=1, _legacy=True):
194
- outputs = model_outputs["logits"][0]
195
- outputs = outputs.numpy()
196
- scores = softmax(outputs)
197
-
198
- if top_k == 1 and _legacy:
199
- return {"label": self.model.config.id2label[scores.argmax().item()], "score": scores.max().item()}
200
-
201
- dict_scores = [
202
- {"label": self.model.config.id2label[i], "score": score.item()} for i, score in enumerate(scores)
203
- ]
204
- if not _legacy:
205
- dict_scores.sort(key=lambda x: x["score"], reverse=True)
206
- if top_k is not None:
207
- dict_scores = dict_scores[:top_k]
208
- return dict_scores
209
-
210
- class EndpointHandler():
211
- def __init__(self, path=""):
212
- # Preload all the elements you are going to need at inference.
213
- PIPELINE_REGISTRY.register_pipeline(
214
- "bibert-multitask-classification",
215
- pipeline_class=BiBert_MultiTaskPipeline,
216
- pt_model=BertForSequenceClassification
217
- )
218
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
219
-
220
- tokenizer = AutoTokenizer.from_pretrained(path)
221
-
222
- tasks = [
223
- Task(id=0, name='label_classification', type='seq_classification', num_labels=5),
224
- Task(id=1, name='binary_classification', type='seq_classification', num_labels=2)
225
- ]
226
-
227
- multitask_model = BertForSequenceClassification.from_pretrained(
228
- path,
229
- tasks_map=tasks
230
- ).to(device)
231
- self.classifier_s = pipeline("bibert-multitask-classification", model = multitask_model, task_id="0", tokenizer=tokenizer, device = device)
232
- self.classifier_p = pipeline("bibert-multitask-classification", model = multitask_model, task_id="1", tokenizer=tokenizer, device = device)
233
-
234
- def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
235
- """
236
- data args:
237
- inputs (:obj: `str` | `PIL.Image` | `np.array`)
238
- kwargs
239
- Return:
240
- A :obj:`list` | `dict`: will be serialized and returned
241
- """
242
- inputs = data.pop("text", data)
243
- lang = data.pop("lang", None)
244
- if isinstance(inputs, str):
245
- inputs = [inputs]
246
-
247
- prediction_p = self.classifier_p(inputs)
248
- print(prediction_p)
249
- label = prediction_p[0]['label']
250
- score = prediction_p[0]['score']
251
-
252
- if label == '0' and score >= 0.75:
253
- label = 2
254
- return {"label":label, "score": score}
255
- else:
256
- prediction_s = self.classifier_s(inputs)
257
- label = prediction_s[0]['label']
258
- score = prediction_s[0]['score']
259
- return prediction_s