# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/napoler/mlmxl-pytorch/blob/master/Memory_Transformer_XL.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + colab={"base_uri": "https://localhost:8080/"} id="L40CAu3TIsAD" outputId="cbc5ea0c-98c3-4e6a-f7db-52377c7a9dfa"
# !Memory Transformer-XL

# + colab={"base_uri": "https://localhost:8080/"} id="-Fl_YYGWKyj4" outputId="01b1bbed-118b-427b-9c0f-3ddbfb51ef5f"
# !pip install memory-transformer-xl
# !pip install transformers
# !pip install mlm-pytorch

# + id="6zyv36DoK75H"
import torch
from memory_transformer_xl import MemoryTransformerXL

model = MemoryTransformerXL(
    num_tokens = 20000,
    dim = 1024,
    heads = 8,
    depth = 8,
    seq_len = 512,
    mem_len = 256,            # short term memory (the memory from transformer-xl)
    lmem_len = 256,           # long term memory (memory attention network attending to short term memory and hidden activations)
    mem_write_iters = 2,      # number of iterations of attention for writing to memory
    memory_layers = [6,7,8],  # which layers to use memory, only the later layers are actually needed
    num_mem_kv = 128,         # number of memory key/values, from All-attention paper

).cuda()

x1 = torch.randint(0, 20000, (1, 512)).cuda()
logits1, mem1 = model(x1)

x2 = torch.randint(0, 20000, (1, 512)).cuda()
logits2, mem2 = model(x2, memories = mem1)

# + colab={"base_uri": "https://localhost:8080/"} id="HJlHpMRELN0F" outputId="20451ae7-b6c1-4595-e2b9-921cdb6baf9d"
mem2

# + id="a5OxkLqPMhb2"


# + [markdown] id="fIgrTMHQMhx8"
# # 分词
# 数据特点：
# 可直接用于预训练、语言模型或语言生成任务。
# 发布专用于简体中文NLP任务的小词表。
# 词表介绍
# Google原始中文词表和我们发布的小词表的统计信息如下：
#
# Token Type	Google	CLUE
# Simplified Chinese	11378	5689
# Traditional Chinese	3264	✗
# English	3529	1320
# Japanese	573	✗
# Korean	84	✗
# Emoji	56	✗
# Numbers	1179	140
# Special Tokens	106	106
# Other Tokens	959	766
# Total	21128	8021
# https://github.com/CLUEbenchmark/CLUEPretrainedModels

# + id="iPNjiqH2Mq6L"
from transformers import AutoTokenizer, AutoModel,BertTokenizer
  
tokenizer = BertTokenizer.from_pretrained("clue/roberta_chinese_clue_tiny")

# model = AutoModel.from_pretrained("clue/roberta_chinese_clue_tiny")

# + colab={"base_uri": "https://localhost:8080/"} id="nl3QCP-GM9bt" outputId="5cd44d38-7204-47a2-b7af-77986d845b66"
tokenizer

# + id="9MybAvPCP_uT"


# + colab={"base_uri": "https://localhost:8080/"} id="7hc4wmOwLLcZ" outputId="d71265f9-48c5-4558-d56c-25ecf21aaa27"
tokenizer.vocab_size

# + colab={"base_uri": "https://localhost:8080/"} id="Evp0oyz5QAKg" outputId="84c395df-c7a8-43f1-f209-2c7743cf3e5a"
dir(tokenizer)

# + [markdown] id="nuaZnZ2wNQKj"
# # 模型测试
#
# 使用mlm模式训练
# https://github.com/lucidrains/mlm-pytorch

# + id="R6-iOR52NUNt"
# vocab_size

# + [markdown] id="qbB84zr1Rb53"
# ## 重新mlm模型
#
#

# + id="0xWjqtquRgi6"
import math
from functools import reduce

import torch
from torch import nn
import torch.nn.functional as F

# helpers

def prob_mask_like(t, prob):
    return torch.zeros_like(t).float().uniform_(0, 1) < prob

def mask_with_tokens(t, token_ids):
    init_no_mask = torch.full_like(t, False, dtype=torch.bool)
    mask = reduce(lambda acc, el: acc | (t == el), token_ids, init_no_mask)
    return mask

def get_mask_subset_with_prob(mask, prob):
    batch, seq_len, device = *mask.shape, mask.device
    max_masked = math.ceil(prob * seq_len)

    num_tokens = mask.sum(dim=-1, keepdim=True)
    mask_excess = (mask.cumsum(dim=-1) > (num_tokens * prob).ceil())
    mask_excess = mask_excess[:, :max_masked]

    rand = torch.rand((batch, seq_len), device=device).masked_fill(~mask, -1e9)
    _, sampled_indices = rand.topk(max_masked, dim=-1)
    sampled_indices = (sampled_indices + 1).masked_fill_(mask_excess, 0)

    new_mask = torch.zeros((batch, seq_len + 1), device=device)
    new_mask.scatter_(-1, sampled_indices, 1)
    return new_mask[:, 1:].bool()

# main class

class MLMXL(nn.Module):
    def __init__(
        self,
        transformer,
        mask_prob = 0.15,
        replace_prob = 0.9,
        num_tokens = None,
        random_token_prob = 0.,
        mask_token_id = 2,
        pad_token_id = 0,
        mask_ignore_token_ids = []):
        super().__init__()

        self.transformer = transformer
        self.mem=None

        # mlm related probabilities
        self.mask_prob = mask_prob
        self.replace_prob = replace_prob

        self.num_tokens = num_tokens
        self.random_token_prob = random_token_prob

        # token ids
        self.pad_token_id = pad_token_id
        self.mask_token_id = mask_token_id
        self.mask_ignore_token_ids = set([*mask_ignore_token_ids, pad_token_id])

    def forward(self, input, **kwargs):
        # do not mask [pad] tokens, or any other tokens in the tokens designated to be excluded ([cls], [sep])
        # also do not include these special tokens in the tokens chosen at random
        no_mask = mask_with_tokens(input, self.mask_ignore_token_ids)
        mask = get_mask_subset_with_prob(~no_mask, self.mask_prob)

        # get mask indices
        mask_indices = torch.nonzero(mask, as_tuple=True)

        # mask input with mask tokens with probability of `replace_prob` (keep tokens the same with probability 1 - replace_prob)
        masked_input = input.clone().detach()

        # if random token probability > 0 for mlm
        if self.random_token_prob > 0:
            assert self.num_tokens is not None, 'num_tokens keyword must be supplied when instantiating MLM if using random token replacement'
            random_token_prob = prob_mask_like(input, self.random_token_prob)
            random_tokens = torch.randint(0, self.num_tokens, input.shape, device=input.device)
            random_no_mask = mask_with_tokens(random_tokens, self.mask_ignore_token_ids)
            random_token_prob &= ~random_no_mask
            random_indices = torch.nonzero(random_token_prob, as_tuple=True)
            masked_input[random_indices] = random_tokens[random_indices]

        # [mask] input
        replace_prob = prob_mask_like(input, self.replace_prob)
        masked_input = masked_input.masked_fill(mask * replace_prob, self.mask_token_id)

        # mask out any tokens to padding tokens that were not originally going to be masked
        labels = input.masked_fill(~mask, self.pad_token_id)
        if self.mem!=None:
        # get generator output and get mlm loss
          logits,self.mem = self.transformer(masked_input, memories = self.mem, **kwargs)
        else:
          logits,self.mem = self.transformer(masked_input, **kwargs)

        mlm_loss = F.cross_entropy(
            logits.transpose(1, 2),
            labels,
            ignore_index = self.pad_token_id
        )

        return mlm_loss



# + id="r-IU2T2UNOw5"
import torch
from memory_transformer_xl import MemoryTransformerXL
import torch
from torch import nn
from torch.optim import Adam
# from mlm_pytorch import MLM

model = MemoryTransformerXL(
    num_tokens = tokenizer.vocab_size,
    dim = 128,
    heads = 8,
    depth = 8,
    seq_len = 1024,
    mem_len = 256,            # short term memory (the memory from transformer-xl)
    lmem_len = 256,           # long term memory (memory attention network attending to short term memory and hidden activations)
    mem_write_iters = 2,      # number of iterations of attention for writing to memory
    memory_layers = [6,7,8],  # which layers to use memory, only the later layers are actually needed
    num_mem_kv = 128,         # number of memory key/values, from All-attention paper

).cuda()

x1 = torch.randint(0, tokenizer.vocab_size, (1, 1024)).cuda()
logits1, mem1 = model(x1)

x2 = torch.randint(0, tokenizer.vocab_size, (1, 1024)).cuda()
logits2, mem2 = model(x2, memories = mem1)

# + id="e5dm_ttEP4cZ"
tokenizer

# + id="h83e2MA6OMLq"
torch.save(model.state_dict(), "model1024.bin")

# + id="GPHPr8PkP5xB"
# plugin the language model into the MLM trainer

trainer = MLMXL(
    model,
    mask_token_id = tokenizer.mask_token_id,          # the token id reserved for masking
    pad_token_id = tokenizer.pad_token_id,           # the token id for padding
    mask_prob = 0.15,           # masking probability for masked language modeling
    replace_prob = 0.90,        # ~10% probability that token will not be masked, but included in loss, as detailed in the epaper
    mask_ignore_token_ids = [tokenizer.cls_token_id,tokenizer.sep_token_id]  # other tokens to exclude from masking, include the [cls] and [sep] here
).cuda()

# optimizer

opt = Adam(trainer.parameters(), lr=3e-4)

# one training step (do this for many steps in a for loop, getting new `data` each time)

data = torch.randint(0, tokenizer.vocab_size, (2, 1024)).cuda()

loss = trainer(data)
loss.backward()
opt.step()
opt.zero_grad()

# after much training, the model should have improved for downstream tasks

# torch.save(transformer, f'./pretrained-model.pt')

# + colab={"base_uri": "https://localhost:8080/"} id="yMCmW8DtSLP-" outputId="0981213a-c488-47f6-e89d-5e26eb77b70e"
loss

# + id="log9_CFzN7Qn"
# dir(model)
model

# + id="8WVOOqD8NkvO"
logits2

StarcoderdataJupyterScriptsDedupFiltered