135 lines
5.3 KiB
Python
135 lines
5.3 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
from typing import Any, Dict
|
|
|
|
import torch
|
|
|
|
from modelscope.utils.constant import ModeKeys
|
|
from .base import OfaBasePreprocessor
|
|
|
|
|
|
class OfaTextClassificationPreprocessor(OfaBasePreprocessor):
|
|
r"""
|
|
OFA preprocessor for text classification tasks.
|
|
"""
|
|
|
|
def __init__(self,
|
|
cfg,
|
|
model_dir,
|
|
mode=ModeKeys.INFERENCE,
|
|
*args,
|
|
**kwargs):
|
|
"""preprocess the data
|
|
|
|
Args:
|
|
cfg(modelscope.utils.config.ConfigDict) : model config
|
|
model_dir (str): model path,
|
|
mode: preprocessor mode (model mode)
|
|
"""
|
|
super(OfaTextClassificationPreprocessor,
|
|
self).__init__(cfg, model_dir, mode, *args, **kwargs)
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
if self.mode == ModeKeys.TRAIN:
|
|
return self._build_train_sample(data)
|
|
else:
|
|
return self._build_infer_sample(data)
|
|
|
|
def _build_instruction(self, data):
|
|
r"""
|
|
Building text classification task's instruction.
|
|
|
|
The `data` should contains key `text` and `text2`, and the final instruction
|
|
is like ` can text1 " {} " imply text2 " {} "?`, the first `{}` refer to
|
|
the value of `text` and the latter refer to `text2`
|
|
|
|
step 1. Preprocess for input text `text` and `text2` in `data`.
|
|
- Do lower, stripe and restrict the maximum length as `max_src_length`.
|
|
step 2. Using instruction template to generate the final instruction.
|
|
step 3. Tokenize the instruction as result.
|
|
"""
|
|
text1 = ' '.join(
|
|
data['text'].lower().strip().split()[:self.max_src_length])
|
|
text2 = ' '.join(
|
|
data['text2'].lower().strip().split()[:self.max_src_length])
|
|
prompt = ' can text1 " {} " imply text2 " {} "?'
|
|
text = prompt.format(text1, text2)
|
|
instruction_itm = self.tokenize_text(text)
|
|
return instruction_itm
|
|
|
|
def _build_train_sample(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
r"""
|
|
Building training samples.
|
|
|
|
step 1. Building instruction for text classification using `_build_instruction`.
|
|
step 2. If the `label` is not text, transfer it to text using `label2ans`.
|
|
step 3. Tokenize the label data.
|
|
step 4. Concatenate the instruction and label tokens as the target item.
|
|
- padding the instruction tokens from target item as `target`.
|
|
- remove the eos token from target item as `prev_output_tokens`.
|
|
step 5. Add constraint mask.
|
|
|
|
Args:
|
|
data (`Dict[str, Any]`): Input data, should contains the key of `text`, `text2`
|
|
and `label`, both of them refer to a text input, and the target of this job
|
|
is to find whether or not `text` imply `text2`, the `label` is the supervised
|
|
data for training.
|
|
Return:
|
|
A dict object, contains source text input, target tokens and previous output
|
|
tokens and constraint mask.
|
|
"""
|
|
instruction_itm = self._build_instruction(data)
|
|
assert 'label' in data, 'there must has `label` column in train phase '
|
|
label = data['label']
|
|
if self.label2ans:
|
|
label = self.label2ans[label] # ans
|
|
label_itm = self.tokenize_text(f' {label}', add_bos=False)
|
|
if self.prompt_type == 'none':
|
|
target_itm = label_itm
|
|
elif self.prompt_type == 'prev_output':
|
|
target_itm = torch.cat([instruction_itm[1:-1], label_itm])
|
|
else:
|
|
raise NotImplementedError
|
|
prev_output_itm = torch.cat([self.bos_item, target_itm[:-1]])
|
|
target_itm[:-len(label_itm)] = self.pad_item
|
|
sample = {
|
|
'source': instruction_itm,
|
|
'target': target_itm,
|
|
'prev_output_tokens': prev_output_itm,
|
|
}
|
|
self.add_constraint_mask(sample)
|
|
return sample
|
|
|
|
def _build_infer_sample(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
r"""
|
|
Building inference samples.
|
|
|
|
step 1. Building instruction for text classification using `_build_instruction`.
|
|
step 2. Whether or not to add `prefix_token`.
|
|
step 3. Whether or not to add `label` data.
|
|
|
|
Args:
|
|
data (`Dict[str, Any]`): Input data, should contains the key of `text` and `text2`,
|
|
both of them refer to a text input, and the target of this job is to find
|
|
whether or not `text` imply `text2`.
|
|
Return:
|
|
A dict object, contains source text input, prefix tokens and label data.
|
|
"""
|
|
instruction_itm = self._build_instruction(data)
|
|
if self.prompt_type == 'none':
|
|
prefix_token = []
|
|
decoder_prompt = self.bos_item
|
|
elif self.prompt_type == 'prev_output':
|
|
prefix_token = instruction_itm[:-1] # remove eos
|
|
decoder_prompt = instruction_itm[:-1]
|
|
else:
|
|
raise NotImplementedError
|
|
sample = {
|
|
'source': instruction_itm,
|
|
'prefix_token': prefix_token,
|
|
'decoder_prompt': decoder_prompt,
|
|
}
|
|
if 'label' in data:
|
|
sample['label'] = self.label2ans[data['label']]
|
|
return sample
|