වර්ගීකරණයසඳහා NLP ආකෘති පුහුණුකරු

11from collections import Counter
12from typing import Callable
13
14import torch
15import torchtext
16from torch import nn
17from torch.utils.data import DataLoader
18import torchtext.vocab
19from torchtext.vocab import Vocab
20
21from labml import lab, tracker, monit
22from labml.configs import option
23from labml_helpers.device import DeviceConfigs
24from labml_helpers.metrics.accuracy import Accuracy
25from labml_helpers.module import Module
26from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex
27from labml_nn.optimizers.configs import OptimizerConfigs

30class NLPClassificationConfigs(TrainValidConfigs):

41    optimizer: torch.optim.Adam

43    device: torch.device = DeviceConfigs()

46    model: Module

48    batch_size: int = 16

50    seq_len: int = 512

52    vocab: Vocab = 'ag_news'

54    n_tokens: int

56    n_classes: int = 'ag_news'

58    tokenizer: Callable = 'character'

61    is_save_models = True

64    loss_func = nn.CrossEntropyLoss()

66    accuracy = Accuracy()

68    d_model: int = 512

70    grad_norm_clip: float = 1.0

73    train_loader: DataLoader = 'ag_news'

75    valid_loader: DataLoader = 'ag_news'

80    is_log_model_params_grads: bool = False

85    is_log_model_activations: bool = False

87    def init(self):

92        tracker.set_scalar("accuracy.*", True)
93        tracker.set_scalar("loss.*", True)

95        hook_model_outputs(self.mode, self.model, 'model')

100        self.state_modules = [self.accuracy]

102    def step(self, batch: any, batch_idx: BatchIndex):

108        data, target = batch[0].to(self.device), batch[1].to(self.device)

111        if self.mode.is_train:
112            tracker.add_global_step(data.shape[1])

115        with self.mode.update(is_log_activations=batch_idx.is_last and self.is_log_model_activations):

119            output, *_ = self.model(data)

122        loss = self.loss_func(output, target)
123        tracker.add("loss.", loss)

126        self.accuracy(output, target)
127        self.accuracy.track()

130        if self.mode.is_train:

132            loss.backward()

134            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)

136            self.optimizer.step()

138            if batch_idx.is_last and self.is_log_model_params_grads:
139                tracker.add('model', self.model)

141            self.optimizer.zero_grad()

144        tracker.save()

147@option(NLPClassificationConfigs.optimizer)
148def _optimizer(c: NLPClassificationConfigs):

153    optimizer = OptimizerConfigs()
154    optimizer.parameters = c.model.parameters()
155    optimizer.optimizer = 'Adam'
156    optimizer.d_model = c.d_model
157
158    return optimizer

161@option(NLPClassificationConfigs.tokenizer)
162def basic_english():

176    from torchtext.data import get_tokenizer
177    return get_tokenizer('basic_english')

180def character_tokenizer(x: str):

184    return list(x)

187@option(NLPClassificationConfigs.tokenizer)
188def character():

192    return character_tokenizer

195@option(NLPClassificationConfigs.n_tokens)
196def _n_tokens(c: NLPClassificationConfigs):

200    return len(c.vocab) + 2

203class CollateFunc:

208    def __init__(self, tokenizer, vocab: Vocab, seq_len: int, padding_token: int, classifier_token: int):

216        self.classifier_token = classifier_token
217        self.padding_token = padding_token
218        self.seq_len = seq_len
219        self.vocab = vocab
220        self.tokenizer = tokenizer

222    def __call__(self, batch):

228        data = torch.full((self.seq_len, len(batch)), self.padding_token, dtype=torch.long)

230        labels = torch.zeros(len(batch), dtype=torch.long)

233        for (i, (_label, _text)) in enumerate(batch):

235            labels[i] = int(_label) - 1

237            _text = [self.vocab[token] for token in self.tokenizer(_text)]

239            _text = _text[:self.seq_len]

241            data[:len(_text), i] = data.new_tensor(_text)

244        data[-1, :] = self.classifier_token

247        return data, labels

250@option([NLPClassificationConfigs.n_classes,
251         NLPClassificationConfigs.vocab,
252         NLPClassificationConfigs.train_loader,
253         NLPClassificationConfigs.valid_loader])
254def ag_news(c: NLPClassificationConfigs):

263    train, valid = torchtext.datasets.AG_NEWS(root=str(lab.get_data_path() / 'ag_news'), split=('train', 'test'))

266    with monit.section('Load data'):
267        from labml_nn.utils import MapStyleDataset

270        train, valid = MapStyleDataset(train), MapStyleDataset(valid)

273    tokenizer = c.tokenizer

276    counter = Counter()

278    for (label, line) in train:
279        counter.update(tokenizer(line))

281    for (label, line) in valid:
282        counter.update(tokenizer(line))

284    vocab = torchtext.vocab.vocab(counter, min_freq=1)

287    train_loader = DataLoader(train, batch_size=c.batch_size, shuffle=True,
288                              collate_fn=CollateFunc(tokenizer, vocab, c.seq_len, len(vocab), len(vocab) + 1))

290    valid_loader = DataLoader(valid, batch_size=c.batch_size, shuffle=True,
291                              collate_fn=CollateFunc(tokenizer, vocab, c.seq_len, len(vocab), len(vocab) + 1))

294    return 4, vocab, train_loader, valid_loader