<jupyter_start><jupyter_text>Reference:    * https://github.com/huggingface/nlp/blob/master/notebooks/Overview.ipynb<jupyter_code>class SST2Dataset(torch.utils.data.Dataset):
    def __init__(self, entries_dict, temperature=1):
        super().__init__()
        self.entries_dict = entries_dict
        self.temperature = temperature
    
    def __len__(self):
        return len(self.entries_dict["label"])
    
    def __getitem__(self, idx):
        return (
            self.entries_dict["input_ids"][idx],
            self.entries_dict["attention_mask"][idx].sum(), #input_lengths
            {
                "label": self.entries_dict["label"][idx], 
                "logits": self.entries_dict["logits"][idx] / self.temperature
            }
        )
train_dict, valid_dict, test_dict = torch.load(str(CACHE_DIR / "distill-dicts-augmented.jbl"))
# Instantiate a PyTorch Dataloader around our dataset
TEMPERATURE = 2.
train_loader = torch.utils.data.DataLoader(SST2Dataset(train_dict, temperature=TEMPERATURE), batch_size=64, shuffle=True)
valid_loader = torch.utils.data.DataLoader(SST2Dataset(valid_dict, temperature=TEMPERATURE), batch_size=64, drop_last=False)
test_loader = torch.utils.data.DataLoader(SST2Dataset(test_dict, temperature=1.), batch_size=64, drop_last=False)
ALPHA = 0
DISTILL_OBJECTIVE = torch.nn.MSELoss()

def cross_entropy(logits, targets):
    targets = F.softmax(targets, dim=-1)
    return -(targets * F.log_softmax(logits, dim=-1)).sum(dim=1).mean()

def distill_loss(logits, targets):
#     distill_part = F.binary_cross_entropy_with_logits(
#         logits[:, 1], targets["logits"][:, 1]
#     )
    distill_part = cross_entropy(
        logits, targets["logits"]
    )
    classification_part = F.cross_entropy(
        logits, targets["label"]
    )
    return ALPHA * classification_part + (1-ALPHA) * distill_part
bert_model = BertForSequenceClassification.from_pretrained(str(CACHE_DIR / "sst2_bert_uncased")).cpu()
bert_model.bert.embeddings.word_embeddings.weight.shape
# Note: apex does not support weight dropping
model = get_sequence_model(
    voc_size=bert_model.bert.embeddings.word_embeddings.weight.shape[0],
    emb_size=bert_model.bert.embeddings.word_embeddings.weight.shape[1],
    pad_idx = 0,
    dropoute = 0,
    rnn_hid = 768,
    rnn_layers = 2,
    bidir = True,
    dropouth = 0.25,
    dropouti = 0.25,
    wdrop = 0,
    unit_type = "gru",
    fcn_layers = [512, 2],
    fcn_dropouts = [0.25, 0.25],
    use_attention = True
)
model
# Copy the embedding weights to the LSTM model
try:
    model.embeddings.encoder.emb.weight.data = bert_model.bert.embeddings.word_embeddings.weight.data
except:
    model.embeddings.encoder.weight.data = bert_model.bert.embeddings.word_embeddings.weight.data
# Freeze the embedding layer
for param in model.embeddings.encoder.parameters():
    param.requires_grad = False
model = model.cuda()
# Use only leaf tensors
parameters = [x for x in model.parameters() if x.is_leaf and x.requires_grad]
del bert_model
optimizer = torch.optim.Adam(parameters, lr=1e-3, betas=(0.8, 0.99))
# optimizer = torch.optim.RMSprop(parameters, lr=0.01)
if APEX_AVAILABLE:
    model, optimizer = amp.initialize(
        model, optimizer, opt_level="O1"
    )
class TransposeCallback(Callback):
    def on_batch_inputs(self, bot, input_tensors, targets):
        input_tensors = [input_tensors[0].transpose(1, 0), input_tensors[1]]
        return input_tensors, targets
class DistillTop1Accuracy(Top1Accuracy):
    def __call__(self, truth, pred):
        truth = truth["label"]
        return super().__call__(truth, pred)
total_steps = len(train_loader) * 5

checkpoints = CheckpointCallback(
    keep_n_checkpoints=1,
    checkpoint_dir=CACHE_DIR / "distill_model_cache/",
    monitor_metric="loss"
)
lr_durations = [
    int(total_steps*0.2),
    int(np.ceil(total_steps*0.8))
]
break_points = [0] + list(np.cumsum(lr_durations))[:-1]
callbacks = [
    MovingAverageStatsTrackerCallback(
        avg_window=len(train_loader) // 8,
        log_interval=len(train_loader) // 10
    ),
    LearningRateSchedulerCallback(
        MultiStageScheduler(
            [
                LinearLR(optimizer, 0.01, lr_durations[0]),
                CosineAnnealingLR(optimizer, lr_durations[1])
            ],
            start_at_epochs=break_points
        )
    ),
    checkpoints,
    TransposeCallback()
]
    
bot = BaseBot(
    log_dir = CACHE_DIR / "distill_logs",
    model=model, 
    train_loader=train_loader,
    valid_loader=valid_loader, 
    clip_grad=10.,
    optimizer=optimizer, echo=True,
    criterion=distill_loss,
    callbacks=callbacks,
    pbar=False, use_tensorboard=False,
    use_amp=APEX_AVAILABLE,
    metrics=(DistillTop1Accuracy(),)
)
print(total_steps)

bot.train(
    total_steps=total_steps,
    checkpoint_interval=len(train_loader) // 2
)
bot.load_model(checkpoints.best_performers[0][1])
checkpoints.remove_checkpoints(keep=0)
bot.eval(valid_loader)
bot.eval(test_loader)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11921850079744817, 'fraction_numerical': 0.0171451355661882, 'mean_word_length': 3.810162991371045, 'pattern_counts': {'":': 2, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '92367', 'license': 'MIT', 'n_tokens_mistral': 1783, 'n_tokens_neox': 1714, 'n_words': 275, 'path': 'notebooks/03-6-flat-gru.ipynb', 'repo_name': 'ceshine/transformer_to_rnn', 'size': '5016'}