2 |
3 | 
4 |
5 |
7 |
8 |
9 | 
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
82 |
96 |
97 | ## Not maintained task
98 | Due to time constraints, the following tasks are temporarily not supported
99 | | | |
100 | |-|-|
101 | | Classification | :label: multi-class and multi-label classification |
102 | | Question Answering | :page_with_curl: extractive qa |
103 | | Question Answering | :radio_button: multiple-choice qa |
104 | | Tagging | :eye_speech_bubble: sequence level tagging / sequence level with crf |
105 | | Self-supervise Learning | :diving_mask: mask language model |
106 |
107 | ## Supplement
108 | - [transformers models list](https://huggingface.co/models): you can find any pretrained models here
109 | - [nlprep](https://github.com/voidful/NLPrep): download and preprocessing data in one line
110 | - [nlp2go](https://github.com/voidful/nlp2go): create demo api as quickly as possible.
111 |
112 |
113 | ## Contributing
114 | Thanks for your interest.There are many ways to contribute to this project. Get started [here](https://github.com/voidful/tfkit/blob/master/CONTRIBUTING.md).
115 |
116 | ## License 
117 |
118 | * [License](https://github.com/voidful/tfkit/blob/master/LICENSE)
119 |
120 | ## Icons reference
121 | Icons modify from Freepik from www.flaticon.com
122 | Icons modify from Nikita Golubev from www.flaticon.com
123 |
--------------------------------------------------------------------------------
/tfkit/task/clas/model.py:
--------------------------------------------------------------------------------
1 | from typing import Dict, List, Any
2 |
3 | import torch
4 | from torch import nn, softmax, sigmoid
5 |
6 | from tfkit.task.clas import Preprocessor
7 | from tfkit.utility.base_model import BaseTFKitModel
8 | from tfkit.utility.constants import DEFAULT_MAXLEN, DEFAULT_DROPOUT
9 | from tfkit.utility.loss import FocalLoss, BCEFocalLoss
10 | from tfkit.utility.predictor import ClassificationPredictor
11 |
12 |
13 | class Model(BaseTFKitModel):
14 | """Multi-class and multi-label classification model."""
15 |
16 | def __init__(self, tokenizer, pretrained, tasks_detail: Dict[str, List[str]],
17 | maxlen: int = DEFAULT_MAXLEN, dropout: float = DEFAULT_DROPOUT, **kwargs):
18 | super().__init__(tokenizer, pretrained, maxlen, **kwargs)
19 |
20 | # Initialize classification-specific components
21 | self.dropout = nn.Dropout(dropout)
22 | self.loss_fct = FocalLoss()
23 | self.loss_fct_mt = BCEFocalLoss()
24 |
25 | # Setup multi-task classification heads
26 | self.tasks = dict()
27 | self.tasks_detail = tasks_detail
28 | self.classifier_list = nn.ModuleList()
29 | for task, labels in tasks_detail.items():
30 | self.classifier_list.append(nn.Linear(self.get_hidden_size(), len(labels)))
31 | self.tasks[task] = len(self.classifier_list) - 1
32 |
33 | self._setup_predictor(ClassificationPredictor, Preprocessor)
34 |
35 | def get_all_task(self):
36 | """
37 | list all classification task
38 | :return: tasks list
39 | """
40 | return list(self.tasks.keys())
41 |
42 | def mean_pooling(self, model_output, attention_mask):
43 | """
44 | Mean Pooling - Take attention mask into account for correct averaging
45 | from https://github.com/UKPLab/sentence-transformers
46 | modify - mask from -1 to 0
47 | :param model_output:
48 | :param attention_mask:
49 | :return:
50 | """
51 | input_mask_expanded = attention_mask.unsqueeze(-1).expand(model_output.size()).float()
52 | input_mask_expanded[input_mask_expanded < 0] = 0
53 | sum_embeddings = torch.sum(model_output * input_mask_expanded, 1)
54 | sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
55 | return sum_embeddings / sum_mask
56 |
57 | def forward(self, batch_data, eval=False, **kwargs):
58 | # covert input to correct data type
59 | tasks = batch_data['task']
60 | tasks = [bytes(t).decode(encoding="utf-8", errors="ignore") for t in tasks]
61 | inputs = torch.as_tensor(batch_data['input'])
62 | targets = torch.as_tensor(batch_data['target'])
63 | masks = torch.as_tensor(batch_data['mask'])
64 | # define model output
65 | result_dict = {
66 | 'max_item': [],
67 | 'prob_list': [],
68 | 'label_prob': []
69 | }
70 |
71 | result_logits = []
72 | result_labels = []
73 | for p, zin in enumerate(zip(tasks, inputs, masks)):
74 | task, input, mask = zin
75 | task_id = self.tasks[task]
76 | task_labels = self.tasks_detail[task]
77 | output = self.pretrained(input.unsqueeze(0), mask.unsqueeze(0))[0]
78 | pooled_output = self.dropout(self.mean_pooling(output, mask.unsqueeze(0)))
79 | classifier_output = self.classifier_list[task_id](pooled_output)
80 | reshaped_logit = classifier_output.view(-1, len(task_labels)) # 0 for cls position
81 | result_logits.append(reshaped_logit)
82 | if not eval:
83 | target = targets[p]
84 | result_labels.append(target)
85 | else:
86 | if 'multi_label' in task:
87 | reshaped_logit = sigmoid(reshaped_logit)
88 | else:
89 | reshaped_logit = softmax(reshaped_logit, dim=1)
90 | logit_prob = reshaped_logit[0].data.tolist()
91 | logit_label = dict(zip(task_labels, logit_prob))
92 | result_dict['label_prob'].append({task: logit_label})
93 | if 'multi_label' in task:
94 | result_dict['max_item'].append({task: [k for k, v in logit_label.items() if v > 0.5]})
95 | else:
96 | result_dict['max_item'].append({task: [task_labels[logit_prob.index(max(logit_prob))]]})
97 |
98 | if eval:
99 | outputs = result_dict
100 | else:
101 | loss = 0
102 | for logit, labels, task in zip(result_logits, result_labels, tasks):
103 | if 'multi_label' in task:
104 | loss += self.loss_fct_mt(logit, labels.type_as(logit))
105 | else:
106 | loss += self.loss_fct(logit, labels)
107 | outputs = loss
108 |
109 | return outputs
110 |
--------------------------------------------------------------------------------
/tests/test_model_loader.py:
--------------------------------------------------------------------------------
1 | from types import SimpleNamespace
2 | from unittest.mock import MagicMock
3 |
4 | import pytest
5 |
6 | from tfkit.utility import model as model_utils
7 | from tfkit.utility.model import load_pretrained_model, load_pretrained_tokenizer
8 |
9 |
10 | def _make_config(**overrides):
11 | defaults = {
12 | "is_encoder_decoder": False,
13 | "architectures": [],
14 | "is_decoder": False,
15 | }
16 | defaults.update(overrides)
17 | return SimpleNamespace(**defaults)
18 |
19 |
20 | def test_load_pretrained_model_prefers_seq2seq(monkeypatch):
21 | config = _make_config(is_encoder_decoder=True)
22 |
23 | auto_config = MagicMock()
24 | auto_config.from_pretrained.return_value = config
25 | monkeypatch.setattr(model_utils, "AutoConfig", auto_config)
26 |
27 | seq2seq_loader = MagicMock()
28 | seq2seq_instance = object()
29 | seq2seq_loader.from_pretrained.return_value = seq2seq_instance
30 | monkeypatch.setattr(model_utils, "AutoModelForSeq2SeqLM", seq2seq_loader)
31 |
32 | causal_loader = MagicMock()
33 | monkeypatch.setattr(model_utils, "AutoModelForCausalLM", causal_loader)
34 |
35 | base_loader = MagicMock()
36 | monkeypatch.setattr(model_utils, "AutoModel", base_loader)
37 |
38 | result = load_pretrained_model("mock-model", ["seq2seq"]) # type: ignore[arg-type]
39 |
40 | assert result is seq2seq_instance
41 | seq2seq_loader.from_pretrained.assert_called_once()
42 | causal_loader.from_pretrained.assert_not_called()
43 | base_loader.from_pretrained.assert_not_called()
44 |
45 |
46 | def test_load_pretrained_model_prefers_causal(monkeypatch):
47 | config = _make_config(architectures=["CustomForCausalLM"])
48 |
49 | auto_config = MagicMock()
50 | auto_config.from_pretrained.return_value = config
51 | monkeypatch.setattr(model_utils, "AutoConfig", auto_config)
52 |
53 | seq2seq_loader = MagicMock()
54 | monkeypatch.setattr(model_utils, "AutoModelForSeq2SeqLM", seq2seq_loader)
55 |
56 | causal_loader = MagicMock()
57 | causal_instance = object()
58 | causal_loader.from_pretrained.return_value = causal_instance
59 | monkeypatch.setattr(model_utils, "AutoModelForCausalLM", causal_loader)
60 |
61 | base_loader = MagicMock()
62 | monkeypatch.setattr(model_utils, "AutoModel", base_loader)
63 |
64 | result = load_pretrained_model("mock-model", ["clm"]) # type: ignore[arg-type]
65 |
66 | assert result is causal_instance
67 | causal_loader.from_pretrained.assert_called_once()
68 | base_loader.from_pretrained.assert_not_called()
69 |
70 |
71 | def test_load_pretrained_model_causal_fallback(monkeypatch):
72 | config = _make_config(architectures=["CustomForCausalLM"])
73 |
74 | auto_config = MagicMock()
75 | auto_config.from_pretrained.return_value = config
76 | monkeypatch.setattr(model_utils, "AutoConfig", auto_config)
77 |
78 | seq2seq_loader = MagicMock()
79 | monkeypatch.setattr(model_utils, "AutoModelForSeq2SeqLM", seq2seq_loader)
80 |
81 | causal_loader = MagicMock()
82 | causal_loader.from_pretrained.side_effect = ValueError("missing head")
83 | monkeypatch.setattr(model_utils, "AutoModelForCausalLM", causal_loader)
84 |
85 | base_loader = MagicMock()
86 | base_instance = object()
87 | base_loader.from_pretrained.return_value = base_instance
88 | monkeypatch.setattr(model_utils, "AutoModel", base_loader)
89 |
90 | result = load_pretrained_model("mock-model", ["clm"]) # type: ignore[arg-type]
91 |
92 | assert result is base_instance
93 | base_loader.from_pretrained.assert_called_once()
94 | assert config.is_decoder is True
95 |
96 |
97 | def test_load_pretrained_model_trust_remote_code_env(monkeypatch):
98 | monkeypatch.setenv("TFKIT_TRUST_REMOTE_CODE", "false")
99 |
100 | config = _make_config()
101 | auto_config = MagicMock()
102 | auto_config.from_pretrained.return_value = config
103 | monkeypatch.setattr(model_utils, "AutoConfig", auto_config)
104 |
105 | base_loader = MagicMock()
106 | base_instance = object()
107 | base_loader.from_pretrained.return_value = base_instance
108 | monkeypatch.setattr(model_utils, "AutoModel", base_loader)
109 |
110 | result = load_pretrained_model("mock-model", ["clas"]) # type: ignore[arg-type]
111 |
112 | assert result is base_instance
113 | auto_config.from_pretrained.assert_called_once_with(
114 | "mock-model", trust_remote_code=False
115 | )
116 | base_loader.from_pretrained.assert_called_once()
117 | _, kwargs = base_loader.from_pretrained.call_args
118 | assert kwargs.get("trust_remote_code") is False
119 |
120 |
121 | def test_load_pretrained_tokenizer_respects_env(monkeypatch):
122 | monkeypatch.setenv("TFKIT_TRUST_REMOTE_CODE", "0")
123 |
124 | tokenizer_loader = MagicMock()
125 | monkeypatch.setattr(model_utils, "AutoTokenizer", tokenizer_loader)
126 |
127 | load_pretrained_tokenizer("mock-tokenizer")
128 |
129 | tokenizer_loader.from_pretrained.assert_called_once_with(
130 | "mock-tokenizer", trust_remote_code=False
131 | )
132 |
--------------------------------------------------------------------------------
/tests/test_task_generation.py:
--------------------------------------------------------------------------------
1 | from types import SimpleNamespace
2 |
3 | import torch
4 | from torch import nn
5 |
6 | from tfkit.task.clm.model import Model as CLMModel
7 | from tfkit.task.seq2seq.model import Model as Seq2SeqModel
8 |
9 |
10 | class DummyTokenizer:
11 | def __init__(self, vocab_size):
12 | self.vocab_size = vocab_size
13 |
14 | def __len__(self):
15 | return self.vocab_size
16 |
17 | def convert_ids_to_tokens(self, idx):
18 | return f"token-{idx}"
19 |
20 |
21 | class DummyCausalPretrained(nn.Module):
22 | def __init__(self):
23 | super().__init__()
24 | self.config = SimpleNamespace(vocab_size=5, hidden_size=4)
25 | self.output_layer = nn.Linear(self.config.hidden_size, self.config.vocab_size)
26 | self.last_kwargs = None
27 |
28 | def get_output_embeddings(self):
29 | return self.output_layer
30 |
31 | def forward(self, input_ids, attention_mask=None, return_dict=True, **kwargs):
32 | self.last_kwargs = kwargs
33 | batch_size, seq_len = input_ids.shape
34 | logits = torch.zeros(batch_size, seq_len, self.config.vocab_size)
35 | outputs = {
36 | "logits": logits,
37 | "last_hidden_state": torch.zeros(batch_size, seq_len, self.config.hidden_size),
38 | }
39 | if "labels" in kwargs:
40 | outputs["loss"] = torch.tensor(0.0)
41 | return outputs
42 |
43 |
44 | class DummyEncoderPretrained(nn.Module):
45 | def __init__(self):
46 | super().__init__()
47 | self.config = SimpleNamespace(vocab_size=5, hidden_size=4)
48 | self.last_kwargs = None
49 |
50 | def get_output_embeddings(self):
51 | return None
52 |
53 | def forward(self, input_ids, attention_mask=None, return_dict=True, **kwargs):
54 | self.last_kwargs = kwargs
55 | batch_size, seq_len = input_ids.shape
56 | hidden = torch.zeros(batch_size, seq_len, self.config.hidden_size)
57 | return {"last_hidden_state": hidden}
58 |
59 |
60 | class DummySeq2SeqPretrained(nn.Module):
61 | def __init__(self):
62 | super().__init__()
63 | self.config = SimpleNamespace(vocab_size=3, hidden_size=4)
64 | self.decoder = nn.Module()
65 | self.output_layer = nn.Linear(self.config.hidden_size, self.config.vocab_size)
66 |
67 | def get_output_embeddings(self):
68 | return self.output_layer
69 |
70 | def forward(
71 | self,
72 | input_ids=None,
73 | attention_mask=None,
74 | decoder_input_ids=None,
75 | decoder_attention_mask=None,
76 | output_hidden_states=False,
77 | use_cache=False,
78 | return_dict=True,
79 | **kwargs,
80 | ):
81 | batch_size, seq_len = decoder_input_ids.shape
82 | hidden = torch.zeros(batch_size, seq_len, self.config.hidden_size)
83 | outputs = {
84 | "last_hidden_state": hidden,
85 | "decoder_hidden_states": (hidden,),
86 | }
87 | return outputs
88 |
89 |
90 | def test_clm_model_uses_pretrained_head_for_loss():
91 | tokenizer = DummyTokenizer(vocab_size=5)
92 | pretrained = DummyCausalPretrained()
93 | model = CLMModel(tokenizer=tokenizer, pretrained=pretrained)
94 |
95 | batch = {
96 | "input": torch.zeros((1, 2), dtype=torch.long),
97 | "mask": torch.ones((1, 2), dtype=torch.long),
98 | "target": torch.tensor([[0, -1]]),
99 | }
100 |
101 | loss = model.forward(batch, eval=False)
102 | assert torch.is_tensor(loss)
103 | assert "labels" in pretrained.last_kwargs
104 | assert pretrained.last_kwargs["labels"].tolist() == [[0, -100]]
105 |
106 | eval_batch = {
107 | **batch,
108 | "start": [0],
109 | }
110 | result = model.forward(eval_batch, eval=True)
111 | assert isinstance(result, dict)
112 | assert "max_item" in result
113 |
114 |
115 | def test_clm_model_falls_back_to_linear_head():
116 | tokenizer = DummyTokenizer(vocab_size=5)
117 | pretrained = DummyEncoderPretrained()
118 | model = CLMModel(tokenizer=tokenizer, pretrained=pretrained)
119 |
120 | batch = {
121 | "input": torch.zeros((1, 2), dtype=torch.long),
122 | "mask": torch.ones((1, 2), dtype=torch.long),
123 | "target": torch.tensor([[0, -1]]),
124 | }
125 |
126 | loss = model.forward(batch, eval=False)
127 | assert torch.is_tensor(loss)
128 | assert pretrained.last_kwargs == {}
129 |
130 |
131 | def test_seq2seq_model_uses_pretrained_output_head():
132 | tokenizer = DummyTokenizer(vocab_size=3)
133 | pretrained = DummySeq2SeqPretrained()
134 | model = Seq2SeqModel(tokenizer=tokenizer, pretrained=pretrained)
135 |
136 | batch = {
137 | "input": torch.zeros((1, 1), dtype=torch.long),
138 | "prev": torch.zeros((1, 1), dtype=torch.long),
139 | "encoder_mask": torch.ones((1, 1), dtype=torch.long),
140 | "decoder_mask": torch.ones((1, 1), dtype=torch.long),
141 | "target": torch.zeros((1, 1), dtype=torch.long),
142 | "ntarget": torch.full((1, 1), -1),
143 | }
144 |
145 | loss = model.forward(batch, eval=False)
146 | assert torch.is_tensor(loss)
147 | assert model.model is pretrained.output_layer
148 |
--------------------------------------------------------------------------------
/tfkit/task/seq2seq/preprocessor.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | import tfkit.utility.tok as tok
4 | from tfkit.utility.data_filereader import get_gen_data_from_file
5 | from tfkit.utility.data_processor import GeneralNLPPreprocessor
6 |
7 |
8 | class Preprocessor(GeneralNLPPreprocessor):
9 | def read_file_to_data(self, path):
10 | return get_gen_data_from_file(path)
11 |
12 | def set_global_parameters(self):
13 | self.tokenize_target = True
14 |
15 | def preprocess_component_convert_to_id(self, item, likelihood=['none', 'pos', 'neg', 'both'], **param_dict):
16 | likelihood = likelihood[0] if isinstance(likelihood, list) else likelihood
17 | tokenized_input, tokenized_target, n_target, b_target = item['input'], \
18 | item.get('target', None), \
19 | item.get('ntarget', None), \
20 | item.get('btarget', None)
21 | previous = item.get("previous", [])
22 | if tokenized_target is None:
23 | yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
24 | 'previous': self.tokenizer.convert_tokens_to_ids(previous)}
25 | elif b_target and len(b_target) > 0:
26 | yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
27 | 'previous': self.tokenizer.convert_tokens_to_ids(previous),
28 | 'target': self.tokenizer.convert_tokens_to_ids(tokenized_target),
29 | 'btarget': self.tokenizer.encode(b_target)}
30 | else:
31 | if "neg" in likelihood or 'both' in likelihood:
32 | # formatting neg data in csv
33 | if n_target is None:
34 | ntext_arr = [
35 | tok.tok_sep(self.tokenizer) + self.tokenizer.convert_tokens_to_string(tokenized_target)]
36 | elif tok.tok_sep(self.tokenizer) in n_target:
37 | ntext_arr = [ntext.strip() for ntext in n_target.split(tok.tok_sep(self.tokenizer))]
38 | else:
39 | ntext_arr = [n_target.strip()]
40 | for neg_text in ntext_arr:
41 | yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
42 | 'previous': self.tokenizer.convert_tokens_to_ids(previous),
43 | 'target': self.tokenizer.convert_tokens_to_ids(tokenized_target),
44 | 'ntarget': self.tokenizer.encode(neg_text)}
45 | else:
46 | yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
47 | 'previous': self.tokenizer.convert_tokens_to_ids(previous),
48 | 'target': self.tokenizer.convert_tokens_to_ids(tokenized_target)}
49 |
50 | # whole sentence masking
51 | if 'pos' in likelihood:
52 | yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
53 | 'target': self.tokenizer.convert_tokens_to_ids(tokenized_target),
54 | 'previous': self.tokenizer.convert_tokens_to_ids(
55 | [tok.tok_mask(self.tokenizer)] * len(tokenized_target))}
56 | elif 'both' in likelihood:
57 | for neg_text in ntext_arr:
58 | yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
59 | 'target': self.tokenizer.convert_tokens_to_ids(tokenized_target),
60 | 'previous': self.tokenizer.convert_tokens_to_ids(
61 | [tok.tok_mask(self.tokenizer)] * len(tokenized_target)),
62 | 'ntarget': self.tokenizer.encode(neg_text)}
63 |
64 | def postprocess(self, item, tokenizer, maxlen, **kwargs):
65 | t_input_id, previous = item['input'], item['previous']
66 | row_dict = {}
67 | if 'target' in item:
68 | target = item['target']
69 | tokenized_target_id = []
70 | if len(previous) == len(target):
71 | tokenized_prev_id = [self.tok_mask_id] * maxlen
72 | else:
73 | tokenized_prev_id = [self.tok_sep_id] + target
74 | tokenized_target_id.extend(target + [self.tok_sep_id])
75 | row_dict['target'] = tokenized_target_id
76 | row_dict['target_pad'] = [-1]
77 | row_dict['prev'] = tokenized_prev_id
78 | row_dict['ntarget'] = [-1] * maxlen
79 | if 'ntarget' in item and len(item['ntarget']) > 0:
80 | tokenized_ntarget_id = item['ntarget']
81 | if len(tokenized_ntarget_id) <= maxlen:
82 | row_dict['ntarget'] = tokenized_ntarget_id
83 | if 'btarget' in item and len(item['btarget']) > 0:
84 | row_dict['btarget'] = tokenizer.encode(item['btarget'])
85 | else:
86 | tokenized_prev_id = [self.tok_sep_id]
87 | tokenized_prev_id.extend(previous)
88 | row_dict['prev'] = tokenized_prev_id
89 |
90 | row_dict['input'] = t_input_id
91 | row_dict['encoder_mask'] = [1] * len(t_input_id)
92 | row_dict['decoder_mask'] = [1] * len(tokenized_prev_id)
93 | return {key: torch.tensor(value) for key, value in row_dict.items()}
94 |
--------------------------------------------------------------------------------
/tfkit/utility/dataset.py:
--------------------------------------------------------------------------------
1 | import os
2 | from collections import defaultdict
3 | from random import choice
4 |
5 | import joblib
6 | import nlp2
7 | from torch.utils import data
8 | from tqdm.contrib.concurrent import process_map
9 |
10 | from tfkit.utility.constants import CACHE_EXTENSION
11 |
12 | try:
13 | from datasets import load_dataset
14 | except Exception: # pragma: no cover - optional dependency
15 | load_dataset = None
16 |
17 |
18 | def get_dataset(file_path, task_class, tokenizer, parameter):
19 | panel = nlp2.Panel()
20 | # all_arg = nlp2.function_get_all_arg_with_value(task_class.preprocessor.prepare_convert_to_id)
21 | # if parameter.get('panel'):
22 | # print("Operation panel for data preprocessing.")
23 | # for missarg in nlp2.function_check_missing_arg(task_class.preprocessor,
24 | # parameter):
25 | # panel.add_element(k=missarg, v=all_arg[missarg], msg=missarg, default=all_arg[missarg])
26 | # filled_arg = panel.get_result_dict()
27 | # parameter.update(filled_arg)
28 | if load_dataset is not None and not os.path.isfile(file_path):
29 | try:
30 | hf_ds = load_dataset(file_path, split=parameter.get('split', 'train'))
31 | return HFDataset(hf_ds, tokenizer=tokenizer,
32 | preprocessor=task_class.Preprocessor,
33 | preprocessing_arg=parameter)
34 | except Exception:
35 | pass
36 | ds = TFKitDataset(fpath=file_path, tokenizer=tokenizer,
37 | preprocessor=task_class.Preprocessor,
38 | preprocessing_arg=parameter)
39 | return ds
40 |
41 |
42 | class TFKitDataset(data.Dataset):
43 | def __init__(self, fpath, tokenizer, preprocessor, preprocessing_arg={}):
44 | cache_path = fpath + "_" + tokenizer.name_or_path.replace("/", "_") + CACHE_EXTENSION
45 | self.task_dict = {}
46 | self.preprocessor = preprocessor(tokenizer, kwargs=preprocessing_arg)
47 | self.tokenizer = tokenizer
48 | if os.path.isfile(cache_path) and preprocessing_arg.get('cache', False):
49 | with open(cache_path, "rb") as fo:
50 | outdata = joblib.load(fo)
51 | sample = outdata['sample']
52 | length = outdata['length']
53 | self.task_dict = outdata['task']
54 | else:
55 | print(f"Start preprocessing...")
56 | sample = defaultdict(list)
57 | length = 0
58 | get_data_item = self.preprocessor.read_file_to_data(fpath)
59 | while True:
60 | try:
61 | for items in process_map(self.preprocessor.preprocess, next(get_data_item),
62 | chunksize=1000):
63 | for i in items:
64 | length += 1
65 | for k, v in i.items():
66 | sample[k].append(v)
67 | print(f"loaded {length} data.")
68 | except StopIteration as e:
69 | tasks = e.value
70 | break
71 | self.task_dict = tasks
72 | print(f"There are {length} datas after preprocessing.")
73 | if preprocessing_arg.get('cache', False):
74 | with open(cache_path, 'wb') as fo:
75 | outdata = {'sample': sample, 'task': self.task_dict, 'length': length}
76 | joblib.dump(outdata, fo)
77 | self.length = length
78 | self.sample = sample
79 | self.task = self.task_dict
80 |
81 | def increase_with_sampling(self, total):
82 | for _ in range(total - self.length):
83 | for key in self.sample.keys():
84 | self.sample[key].append(choice(self.sample[key]))
85 |
86 | def __len__(self):
87 | return self.length
88 |
89 | def __getitem__(self, idx):
90 | return self.preprocessor.postprocess(
91 | {**{'task_dict': self.task_dict}, **{key: self.sample[key][idx] for key in self.sample.keys()}},
92 | self.tokenizer,
93 | maxlen=self.preprocessor.parameters['maxlen'])
94 |
95 |
96 | class HFDataset(data.Dataset):
97 | """Dataset wrapper for the HuggingFace datasets library."""
98 |
99 | def __init__(self, hf_dataset, tokenizer, preprocessor, preprocessing_arg=None):
100 | preprocessing_arg = preprocessing_arg or {}
101 | self.task_dict = {}
102 | self.sample = defaultdict(list)
103 | self.preprocessor = preprocessor(tokenizer, kwargs=preprocessing_arg)
104 | self.tokenizer = tokenizer
105 |
106 | print("Start preprocessing with HuggingFace dataset...")
107 | length = 0
108 | for raw_item in hf_dataset:
109 | for items in self.preprocessor.preprocess(raw_item):
110 | length += 1
111 | for k, v in items.items():
112 | self.sample[k].append(v)
113 | self.length = length
114 | self.task = self.task_dict
115 |
116 | def __len__(self):
117 | return self.length
118 |
119 | def __getitem__(self, idx):
120 | return self.preprocessor.postprocess(
121 | {**{'task_dict': self.task_dict}, **{key: self.sample[key][idx] for key in self.sample.keys()}},
122 | self.tokenizer,
123 | maxlen=self.preprocessor.parameters['maxlen'])
124 |
--------------------------------------------------------------------------------
/tfkit/test/utility/test_utility_loss.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 |
4 | import torch
5 | from torch import nn
6 | from torch.autograd import Variable
7 |
8 | dir_path = os.path.dirname(os.path.realpath(__file__))
9 | sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
10 |
11 | import unittest
12 | import tfkit
13 |
14 |
15 | class TestLoss(unittest.TestCase):
16 | outputs = Variable(torch.Tensor([[0.00000000000009, 5, 0.5], [0.00000000000000000001, 69, 9]]), requires_grad=False)
17 | targets = Variable(torch.Tensor([1, 1]).long(), requires_grad=False)
18 | alln_targets = Variable(torch.Tensor([-1, -1]).long(), requires_grad=False)
19 | onen_targets = Variable(torch.Tensor([1, -1]).long(), requires_grad=False)
20 |
21 | def testLabelSmoothingCrossEntropy(self):
22 | outputs = torch.Tensor([[0.00000000000009, 5, 0.5], [0.00000000000000000001, 69, 9]])
23 | targets = torch.Tensor([1, 1]).long()
24 | alln_targets = torch.Tensor([0, -1]).long()
25 | onen_targets = torch.Tensor([1, -1]).long()
26 |
27 | criterion = nn.CrossEntropyLoss(ignore_index=-1)
28 | custom_criterion = tfkit.utility.loss.LabelSmoothingLoss(3, ignore_index=-1)
29 |
30 | self.assertTrue(criterion(outputs, targets).item() <
31 | custom_criterion(outputs, targets).item())
32 | self.assertTrue(criterion(outputs, onen_targets).item() <
33 | custom_criterion(outputs, onen_targets).item())
34 |
35 | criterion = nn.CrossEntropyLoss()
36 | custom_criterion = tfkit.utility.loss.LabelSmoothingLoss(3)
37 | self.assertTrue(criterion(outputs, targets).item() <
38 | custom_criterion(outputs, targets).item())
39 |
40 | custom_criterion = tfkit.utility.loss.LabelSmoothingLoss(3, reduction='none')
41 | print(custom_criterion(self.outputs, self.targets))
42 | self.assertTrue(list(custom_criterion(self.outputs, self.targets).shape) == [2])
43 |
44 | def testDiceLoss(self):
45 | custom_criterion = tfkit.utility.loss.DiceLoss(ignore_index=-1)
46 | self.assertTrue(0.8 < custom_criterion(self.outputs, self.targets).item() < 1)
47 | self.assertTrue(0.99 < custom_criterion(self.outputs, self.alln_targets).item() <= 1)
48 | self.assertTrue(0.8 < custom_criterion(self.outputs, self.onen_targets).item() < 1)
49 |
50 | custom_criterion = tfkit.utility.loss.DiceLoss(reduction='none')
51 | print(custom_criterion(self.outputs, self.targets))
52 | self.assertTrue(list(custom_criterion(self.outputs, self.targets).shape) == [2])
53 |
54 | def testLossDrop(self):
55 | outputs = torch.Tensor([[0.00000000000009, 5, 0.5], [0.00000000000000000001, 69, 9]])
56 | targets = torch.Tensor([1, 1]).long()
57 | norm_loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
58 | loss_fct = nn.CrossEntropyLoss(reduction='none', ignore_index=-1) # -1 index = padding token
59 | masked_lm_loss = loss_fct(outputs, targets)
60 | masked_lm_loss = masked_lm_loss.view(-1, len(targets)) # view by batch size
61 | masked_lm_loss = masked_lm_loss.sum(dim=0)
62 | masked_lm_loss = masked_lm_loss.mean()
63 | print(masked_lm_loss.mean(), norm_loss_fct(outputs, targets).mean())
64 |
65 | def testBCEFocalLoss(self):
66 | outputs = torch.Tensor([[0, 1, 0], [0.2, 0, 0]])
67 | targets = torch.Tensor([[0, 1, 0], [1, 0, 0]])
68 | criterion = nn.BCELoss()
69 | custom_criterion = tfkit.utility.loss.BCEFocalLoss()
70 | self.assertTrue(criterion(outputs, targets).item() >
71 | custom_criterion(outputs, targets).item())
72 |
73 | def testNegativeCElLoss(self):
74 | outputs = torch.Tensor([[0.00000000000009, 5, 0.5], [0.00000000000000000001, 69, 9]])
75 | targets = torch.Tensor([1, 1]).long()
76 | alln_targets = torch.Tensor([-1, -1]).long()
77 | onen_targets = torch.Tensor([1, -1]).long()
78 |
79 | criterion = nn.CrossEntropyLoss(ignore_index=-1)
80 | custom_criterion = tfkit.utility.loss.NegativeCElLoss()
81 | self.assertTrue(
82 | criterion(outputs, targets).item() < custom_criterion(outputs, self.targets).item())
83 | self.assertTrue(criterion(outputs, onen_targets).item() < custom_criterion(outputs, onen_targets).item())
84 |
85 | def testFocalLoss(self):
86 | criterion = nn.CrossEntropyLoss(ignore_index=-1)
87 | custom_criterion = tfkit.utility.loss.FocalLoss(gamma=0)
88 | self.assertAlmostEqual(criterion(self.outputs, self.targets).item(),
89 | custom_criterion(self.outputs, self.targets).item())
90 | self.assertAlmostEqual(criterion(self.outputs, self.alln_targets).item(),
91 | custom_criterion(self.outputs, self.alln_targets).item())
92 | self.assertAlmostEqual(criterion(self.outputs, self.onen_targets).item(),
93 | custom_criterion(self.outputs, self.onen_targets).item())
94 |
95 | custom_criterion = tfkit.utility.loss.FocalLoss(gamma=1)
96 | self.assertTrue(criterion(self.outputs, self.targets) > custom_criterion(self.outputs, self.targets))
97 | self.assertTrue(criterion(self.outputs, self.alln_targets).item() - custom_criterion(self.outputs,
98 | self.alln_targets).item() < 1)
99 | self.assertTrue(criterion(self.outputs, self.onen_targets) > custom_criterion(self.outputs, self.onen_targets))
100 |
--------------------------------------------------------------------------------
/tfkit/utility/loss.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch import nn
3 | import torch.nn.functional as F
4 | from torch.autograd import Variable
5 |
6 |
7 | class BCEFocalLoss(nn.Module):
8 | def __init__(self, gamma=2):
9 | super(BCEFocalLoss, self).__init__()
10 | self.gamma = gamma
11 |
12 | def forward(self, input, target):
13 | BCE_loss = F.binary_cross_entropy_with_logits(input, target, reduction='none')
14 | pt = torch.exp(-BCE_loss) # prevents nans when probability 0
15 | focal_loss = (1 - pt) ** self.gamma * BCE_loss
16 | return focal_loss.mean()
17 |
18 |
19 | class FocalLoss(nn.Module):
20 | def __init__(self, gamma=2, ignore_index=-1):
21 | super(FocalLoss, self).__init__()
22 | self.gamma = gamma
23 | self.softmax = nn.Softmax(dim=1)
24 | self.nll = nn.NLLLoss(ignore_index=ignore_index)
25 |
26 | def forward(self, input, target):
27 | softmax = self.softmax(input)
28 | logpt = torch.log(softmax)
29 | pt = Variable(logpt.data.exp())
30 | return self.nll((1 - pt) ** self.gamma * logpt, target)
31 |
32 |
33 | class SeqCTCLoss(nn.Module):
34 | def __init__(self, blank_index):
35 | super(SeqCTCLoss, self).__init__()
36 | self.blank_index = blank_index
37 |
38 | def forward(self, logits, input_lengths, targets, target_lengths):
39 | # lengths : (batch_size, )
40 | # log_logits : (T, batch_size, n_class), this kind of shape is required for ctc_loss
41 | # log_logits = logits + (logit_mask.unsqueeze(-1) + 1e-45).log()
42 | log_logits = logits.log_softmax(-1).transpose(0, 1)
43 | loss = F.ctc_loss(log_logits,
44 | targets,
45 | input_lengths,
46 | target_lengths,
47 | blank=self.blank_index,
48 | reduction='mean',
49 | zero_infinity=True)
50 | return loss
51 |
52 |
53 | class SelfKDLoss(nn.Module):
54 |
55 | def __init__(self, alpha=0.1, temperature=2,ignore_index=-1):
56 | super(SelfKDLoss, self).__init__()
57 | self.alpha = alpha
58 | self.temperature = temperature
59 | self.ignore_index = ignore_index
60 |
61 | def forward(self, outputs, teacher_outputs, labels):
62 | loss = nn.KLDivLoss()(F.log_softmax(outputs / self.temperature, dim=-1),
63 | F.softmax(teacher_outputs / self.temperature, dim=-1)) * (
64 | self.alpha * self.temperature * self.temperature) + F.cross_entropy(outputs, labels,ignore_index=self.ignore_index,) * (
65 | 1. - self.alpha)
66 | return loss
67 |
68 |
69 | class DiceLoss(nn.Module):
70 | """From 'Dice Loss for Data-imbalanced NLP Tasks'"""
71 |
72 | def __init__(self, ignore_index=None, reduction='mean'):
73 | super(DiceLoss, self).__init__()
74 | self.ignore_index = ignore_index
75 | self.reduction = reduction
76 |
77 | def forward(self, y_pred, y_true):
78 | y_pred = torch.softmax(y_pred, dim=1)
79 | if self.ignore_index is not None:
80 | mask = y_true == -1
81 | filtered_target = y_true
82 | filtered_target[mask] = 0
83 | torch.gather(y_pred, dim=1, index=filtered_target.unsqueeze(1))
84 | mask = mask.unsqueeze(1).expand(y_pred.data.size())
85 | y_pred[mask] = 0
86 | pred_prob = torch.gather(y_pred, dim=1, index=y_true.unsqueeze(1))
87 | dsc_i = 1 - ((1 - pred_prob) * pred_prob) / ((1 - pred_prob) * pred_prob + 1)
88 | if self.reduction == 'mean':
89 | return dsc_i.mean()
90 | else:
91 | return dsc_i.view(-1)
92 |
93 |
94 | class NegativeCElLoss(nn.Module):
95 | def __init__(self, ignore_index=-1, reduction='mean'):
96 | super(NegativeCElLoss, self).__init__()
97 | self.softmax = nn.Softmax(dim=1)
98 | self.alpha = 1
99 | self.nll = nn.NLLLoss(ignore_index=ignore_index, reduction=reduction)
100 |
101 | def forward(self, input, target):
102 | nsoftmax = self.softmax(input)
103 | nsoftmax = torch.clamp((1.0 - nsoftmax), min=1e-32)
104 | return self.nll(torch.log(nsoftmax) * self.alpha, target)
105 |
106 |
107 | class LabelSmoothingLoss(nn.Module):
108 | def __init__(self, classes, smoothing=0.1, dim=-1, ignore_index=None, reduction='mean'):
109 | super(LabelSmoothingLoss, self).__init__()
110 | self.confidence = 1.0 - smoothing
111 | self.smoothing = smoothing
112 | self.cls = classes
113 | self.dim = dim
114 | self.reduction = reduction
115 | self.ignore_index = ignore_index
116 |
117 | def forward(self, pred, target):
118 | pred = pred.log_softmax(dim=self.dim)
119 | with torch.no_grad():
120 | true_dist = torch.zeros_like(pred)
121 | true_dist.fill_(self.smoothing / (self.cls - 1))
122 | if self.ignore_index is not None:
123 | mask = target == -1
124 | filtered_target = target.clone()
125 | filtered_target[mask] = 0
126 | true_dist.scatter_(1, filtered_target.unsqueeze(1), self.confidence)
127 | mask = mask.unsqueeze(1).expand(pred.data.size())
128 | true_dist[mask] = 0
129 | else:
130 | true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
131 | if self.reduction == 'mean':
132 | return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))
133 | else:
134 | return torch.sum(-true_dist * pred, dim=self.dim)
135 |
--------------------------------------------------------------------------------
/tfkit/task/oncectc/model.py:
--------------------------------------------------------------------------------
1 | from collections import defaultdict
2 |
3 | import torch
4 | from torch import nn
5 | from torch.nn.functional import softmax
6 |
7 | from tfkit.task.once import Preprocessor
8 | from tfkit.utility.base_model import BaseTFKitModel
9 | from tfkit.utility.constants import BLANK_TOKEN
10 | from tfkit.utility.loss import *
11 | from tfkit.utility.loss import SeqCTCLoss
12 | from tfkit.utility.predictor import NonAutoRegressivePredictor
13 | from tfkit.utility.tok import *
14 |
15 |
16 | class Model(BaseTFKitModel):
17 | """Once generation model with CTC loss for non-autoregressive text generation."""
18 |
19 | def __init__(self, tokenizer, pretrained, maxlen=512, tasks_detail=None, **kwargs):
20 | super().__init__(tokenizer, pretrained, maxlen, **kwargs)
21 |
22 | # Setup CTC-specific components
23 | self.blank_token = BLANK_TOKEN
24 | self.tokenizer.add_tokens(self.blank_token)
25 | self.pretrained.resize_token_embeddings(len(tokenizer))
26 | self.blank_index = self.tokenizer.convert_tokens_to_ids([self.blank_token])[0]
27 | self.loss = SeqCTCLoss(blank_index=self.blank_index)
28 |
29 | # Update vocab size after adding tokens
30 | self.vocab_size = max(self.pretrained.config.vocab_size, self.tokenizer.__len__())
31 | self.model = nn.Linear(self.get_hidden_size(), self.vocab_size)
32 |
33 | self._setup_predictor(NonAutoRegressivePredictor, Preprocessor)
34 |
35 | def forward(self, batch_data, eval=False, max_return=1, **kwargs):
36 | inputs = batch_data['input']
37 | masks = batch_data['mask']
38 | starts = batch_data['start']
39 | ends = batch_data['end']
40 | tokens_tensor = torch.as_tensor(inputs)
41 | mask_tensors = torch.as_tensor(masks)
42 |
43 | output = self.pretrained(tokens_tensor, attention_mask=mask_tensors)
44 | sequence_output = output[0]
45 | prediction_scores = self.model(sequence_output)
46 | batch_size = list(tokens_tensor.shape)[0]
47 | prediction_scores = prediction_scores.view(batch_size, -1, self.vocab_size)
48 |
49 | if eval:
50 | result_dict = {
51 | 'max_item': [],
52 | 'label_prob': defaultdict(list),
53 | 'prob_list': []
54 | }
55 | start = batch_data['start'][0]
56 | topK_ids = [[]] * max_return
57 | topK_probs = [1] * max_return
58 |
59 | pscore = prediction_scores.detach().cpu()
60 | predicted_indexs = pscore.argmax(2).tolist()[0]
61 | predicted_tokens = self.tokenizer.convert_ids_to_tokens(predicted_indexs)
62 | output = []
63 | for pos, (predicted_index, predicted_token) in enumerate(zip(predicted_indexs, predicted_tokens)):
64 | if len(output) > 0 and predicted_index == output[-1]:
65 | continue
66 | if predicted_token == self.blank_token:
67 | continue
68 | if predicted_token == tok_pad(self.tokenizer):
69 | continue
70 | if predicted_token == tok_sep(self.tokenizer):
71 | break
72 |
73 | softmax_score = softmax(prediction_scores[0][pos], dim=0)
74 | max_item_id = torch.argmax(softmax_score, -1).item()
75 | max_item_prob = softmax_score[max_item_id].item()
76 | if max_return > 1:
77 | topK = torch.topk(softmax_score, max_return)
78 | for k, (prob, tid) in enumerate(zip(topK.values.data.tolist(), topK.indices.data.tolist())):
79 | topK_ids[k].append(tid)
80 | topK_probs[k] *= prob
81 | else:
82 | topK_ids[0].append(max_item_id)
83 | topK_probs[0] *= max_item_prob
84 | start += 1
85 |
86 | result_dict['prob_list'] = topK_probs
87 | result_dict['label_prob'] = [[self.tokenizer.decode(ids), prob] for ids, prob in
88 | zip(topK_ids, topK_probs)]
89 | result_dict['max_item'] = [i[0] for i in result_dict['label_prob']]
90 | outputs = result_dict
91 | else:
92 | targets = batch_data['target']
93 | negative_targets = batch_data['ntarget']
94 | input_lengths = batch_data['input_length']
95 | target_lengths = batch_data['target_length']
96 |
97 | target_tensors = torch.as_tensor(targets)
98 | input_length_tensors = torch.as_tensor(input_lengths)
99 | target_length_tensors = torch.as_tensor(target_lengths)
100 |
101 | loss_tensors = torch.as_tensor(targets)
102 | negativeloss_tensors = torch.as_tensor(negative_targets)
103 | ctc_lm_loss = self.loss(prediction_scores,
104 | input_length_tensors,
105 | target_tensors.view(batch_size, -1),
106 | target_length_tensors)
107 |
108 | loss_fct = nn.CrossEntropyLoss(ignore_index=-1) # -1 index = padding token
109 | masked_lm_loss = loss_fct(prediction_scores.view(-1, self.vocab_size),
110 | loss_tensors.view(-1))
111 | if not torch.all(negativeloss_tensors.eq(-1)).item():
112 | negative_loss_fct = NegativeCElLoss()
113 | negative_loss = negative_loss_fct(prediction_scores.view(-1, self.vocab_size),
114 | negativeloss_tensors.view(-1))
115 | masked_lm_loss += negative_loss
116 | outputs = ctc_lm_loss + masked_lm_loss
117 |
118 | return outputs
119 |
--------------------------------------------------------------------------------
/docs/index.md:
--------------------------------------------------------------------------------
1 | Multi-task training
83 | 84 | ```bash 85 | tfkit-train \ 86 | --task clas clas \ 87 | --config xlm-roberta-base \ 88 | --train training_data_taskA.csv training_data_taskB.csv \ 89 | --test testing_data_taskA.csv testing_data_taskB.csv \ 90 | --lr 4e-5 \ 91 | --maxlen 384 \ 92 | --epoch 10 \ 93 | --savedir roberta_sentiment_classifier_multi_task 94 | ``` 95 |
2 |
3 |
4 |
5 |
7 |
8 |
9 | 