CS336 Assignment 1: Tokenizer & Transformer

def train_bpe(input_path: str, vocab_size: int, special_tokens: list[str]):
    PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""

    vocab = {}
    merges = []
    
    
    # 1. init vocab
    for i in range(256):
        bt = bytes([i])
        vocab[i] = bt
    for tok in special_tokens:
        # special_token 不可分割
        vocab[len(vocab)] = tok.encode("utf-8")
    
    # 2. read train data
    with open(input_path, "r", encoding="utf-8") as f:
        text = f.read()
    
    subwords = []
    
    # 3. split text into subwords
    if special_tokens:
        toks = sorted(special_tokens, key=len, reverse=True)
        union = "|".join(re.escape(t) for t in toks)
        parts = re.split(f"({union})", text)

        st = set(special_tokens)
        for part in parts:
            if not part or part in st:
                continue
            subwords.extend(re.findall(PAT, part))
    else:
        subwords = re.findall(PAT, text)
    
    # 4. count up pairs
    unicodes = [subword.encode("utf-8") for subword in subwords]
    
    while len(vocab) < vocab_size:
        # find pair with max count
        to_merge = (None, 0)
        pair_counts = defaultdict(int)
        for unicode in unicodes:
            for idx in range(len(unicode) - 1):
                pair = (bytes([unicode[idx]]), bytes([unicode[idx+1]]))
                pair_counts[pair] += 1
                
                if (to_merge[0] is None)  or \
                    (pair_counts[pair] > to_merge[1]) or \
                    (pair_counts[pair] == to_merge[1] and pair > to_merge[0]):
                    to_merge = (pair, pair_counts[pair])
        
        if to_merge[0] is None:
            raise
        
        # record
        new_unicodes = []
        new_token = to_merge[0][0] + to_merge[0][1]
        vocab[len(vocab)] = new_token
        merges.append(to_merge[0])
        
        # replace pair
        for idx, unicode in enumerate(unicodes):
            to_replace = [i for i in range(len(unicode)-1) if unicode[i:i+2] == to_merge[0]]
            if not to_replace:
                continue
            i = 0
            new_unicode = []
            while i < len(unicode):
                if i in to_replace:
                    new_unicode.append(new_token)
                    i += 2
                else:
                    new_unicode.append(unicode[i])
                    i += 1
            new_unicodes.append(new_unicode) 

        unicodes = new_unicodes
        
    return vocab, merges

    heap: list[HeapItem] = []
    linkedlist = LinkedList[int]()
    pair_counts: dict[tuple[int, int], int] = defaultdict(int)
    back_link: dict[int, set[LinkNode[int]]] = defaultdict(set)
    
    subwords = [list(subword.encode("utf-8")) for subword in subwords]
    for subword in subwords:
        for idx in range(len(subword) - 1):
            cnt, nxt = subword[idx], subword[idx+1]
            node = linkedlist.push_back(cnt)
            back_link[cnt].add(node)
            pair_counts[(cnt, nxt)] += 1
        linkedlist.push_back(subword[-1])
        linkedlist.push_back(-1)  # 插入一个分割符，来区分不同的 subword
    
    for (a, b), cnt in pair_counts.items():
        item = HeapItem(cnt, a, b, vocab)
        heapq.heappush(heap, item)
    
    while heap and len(vocab) < vocab_size:
        heapitem = heapq.heappop(heap)
        pair = heapitem.get_pair()
        
        # 采用懒删除，所以每次判断是不是存在的pair
        if pair not in pair_counts or pair_counts[pair] != heapitem.count:
            continue

        merges.append((vocab[pair[0]], vocab[pair[1]]))
        new_token = len(vocab)
        vocab[new_token] = vocab[pair[0]] + vocab[pair[1]]
        
        # 修改双向链表，反向链接和pair_counts字典
        related_nodes = list(back_link[pair[0]])
        # --- 1 ---
        pair_counts.pop(pair)
        for node in related_nodes:
            if node.nxt is None or node.value != pair[0] or node.nxt.value != pair[1]:
                continue
            a, b = node, node.nxt
            x, y = a.pre, b.nxt
            
            # --- 2 ---
            back_link[pair[0]].discard(a)
            back_link[pair[1]].discard(b)
            # --- 3 ---
            if x is not None and x.value != -1:
                xa = (x.value, a.value)
                xc = (x.value, new_token)
                pair_counts[xa] -= 1
                pair_counts[xc] += 1
                heapq.heappush(heap, HeapItem(pair_counts[xa], *xa, vocab))
                heapq.heappush(heap, HeapItem(pair_counts[xc], *xc, vocab))
            # --- 4 ---
            if y is not None and y.value != -1:
                by = (b.value, y.value)
                cy = (new_token, y.value)
                pair_counts[by] -= 1
                pair_counts[cy] += 1
                heapq.heappush(heap, HeapItem(pair_counts[by], *by, vocab))
                heapq.heappush(heap, HeapItem(pair_counts[cy], *cy, vocab))
            # --- 5 ---
            a.value = new_token
            linkedlist.delete_node(b)
            # --- 6 ---
            back_link[new_token].add(a)
    
    return vocab, merges

class BPE_Tokenizer:
    
    def __init__(
        self,
        vocab: dict[int, bytes],
        merges: list[tuple[bytes, bytes]],
        special_tokens: list[str] | None = None,
    ):
        self.vocab = vocab.copy()
        self.special_tokens = sorted(special_tokens, key=len, reverse=True)
        self.merges = merges
        self.rev_vocab = {v:k for k,v in vocab.items()}
        self.PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
        self.sp_union = "|".join(re.escape(t) for t in self.special_tokens)
      
    
    @classmethod
    def from_files(
        cls,
        vocab_filepath: str,
        merges_filepath: str,
        special_tokens: list[str] | None = None
    ):

        # Load and normalize vocab: keys -> int, values -> bytes
        with open(vocab_filepath, "r", encoding="utf-8") as vf:
            raw_vocab = json.load(vf)

        norm_vocab = {v:k.encode("utf-8") for k, v in raw_vocab.items()}
        norm_merges = []
        # Load and normalize merges: ensure tuples of bytes
        for line in open(merges_filepath, "rb"):
            a, b = line.split()
            
            norm_merges.append(
                (
                    a.encode("utf-8") if isinstance(a, str) else a, 
                    b.encode("utf-8") if isinstance(b, str) else b
                )
            )

        return cls(norm_vocab, norm_merges, special_tokens)
     
     
    def encode(self, text: str) -> list[int]:
        subwords = []
        if self.special_tokens:
            subwords = re.split(f"({self.sp_union})", text)
        else:
            subwords = re.findall(self.PAT, text)
        
        ids = []
        for subword in subwords:
            ids.extend(self.merge_bytes([self.rev_vocab[bytes([byte_id])] for byte_id in subword.encode("utf-8")]))
        
        return ids
    
    def merge_bytes(self, raw_vocab_ids: list[int]) -> list[int]:
        merged_vocab_ids = []
        # 构建双向链表
        linkedlist = LinkedList[int]()
        # 构建一个反向链接
        back_link = defaultdict(list)
        for raw_vocab_id in raw_vocab_ids:
            node = linkedlist.push_back(raw_vocab_id)
            back_link[raw_vocab_id].append(node)
        
        # 遍历merges
        for merge in self.merges:
            new_token = merge[0] + merge[1]
            for node in list(back_link[self.rev_vocab[merge[0]]]):
                # 假设有 a → b
                a, b = node, node.nxt
                if b is None or self.vocab[b.value] != merge[1]:
                    continue
                
                # a → b 变成 c
                a.value = self.rev_vocab[new_token]
                linkedlist.delete_node(b)
                
                # 更新backlink
                back_link[self.rev_vocab[merge[0]]].remove(a)
                back_link[self.rev_vocab[merge[1]]].remove(b)
                back_link[self.rev_vocab[new_token]].append(a)
              
        head = linkedlist.head  
        while head is not None:
            merged_vocab_ids.append(head.value)
            head = head.nxt
        
        return merged_vocab_ids
             
            
    def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]:
        for chunk in iterable:
            yield from self.encode(chunk)
    
    
    def decode(self, ids: list[int]) -> str:
        return (b''.join([self.vocab[_id] for _id in ids])).decode("utf-8", errors="replace")

def gpt2_bytes_to_unicode() -> dict[int, str]:
    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
    cs = bs[:]
    n = 0
    for b in range(2**8):
        if b not in bs:
            bs.append(b)
            cs.append(2**8 + n)
            n += 1
    characters = [chr(n) for n in cs]
    d = dict(zip(bs, characters))
    return d
    
      
class BPE_Tokenizer:
    @classmethod
    def from_files(
        cls,
        vocab_filepath: str,
        merges_filepath: str,
        special_tokens: list[str] | None = None
    ):
        gpt2_byte_decoder = {v: k for k, v in gpt2_bytes_to_unicode().items()}
        with open(vocab_filepath, encoding="utf-8") as vocab_f:
            gpt2_vocab = json.load(vocab_f)
        gpt2_bpe_merges = []
        with open(merges_filepath, encoding="utf-8") as f:
            for line in f:
                cleaned_line = line.rstrip()
                if cleaned_line and len(cleaned_line.split(" ")) == 2:
                    gpt2_bpe_merges.append(tuple(cleaned_line.split(" ")))
        vocab = {
            gpt2_vocab_index: bytes([gpt2_byte_decoder[token] for token in gpt2_vocab_item])
            for gpt2_vocab_item, gpt2_vocab_index in gpt2_vocab.items()
        }
        if special_tokens:
            for special_token in special_tokens:
                byte_encoded_special_token = special_token.encode("utf-8")
                if byte_encoded_special_token not in set(vocab.values()):
                    vocab[len(vocab)] = byte_encoded_special_token

        merges = [
            (
                bytes([gpt2_byte_decoder[token] for token in merge_token_1]),
                bytes([gpt2_byte_decoder[token] for token in merge_token_2]),
            )
            for merge_token_1, merge_token_2 in gpt2_bpe_merges
        ]
        
        return cls(vocab, merges, special_tokens)
     

    def encode(self, text: str) -> list[int]:
        subwords = []
        out = []
        if self.special_tokens:
            for subword in re.split(f"({self.sp_union})", text):
                subwords.extend([subword] if subword in self.special_tokens else re.findall(self.PAT, subword))
        else:
            subwords = re.findall(self.PAT, text)
        
        for subword in subwords:
            if subword in self.special_tokens:
                out.append(self.rev_vocab[subword.encode("utf-8")])
            else:
                out.extend(self.encode_non_special(subword))
        return out
    
    def encode_non_special(self, subword: str) -> list[int]:
        heap = []
        # 1. subword 变成 vocab-ids
        raw_vocab_ids = [self.rev_vocab[bytes([byte_id])] for byte_id in subword.encode("utf-8")]
        # 2. 用双向链表存文本
        linkedlist = LinkedList[LinkItem]()
        for raw_vocab_id in raw_vocab_ids:
            linkedlist.push_back(LinkItem(vocab_id=raw_vocab_id))
        
        # 小根堆，(rank, node)
        node = linkedlist.head
        while node is not None:
            rank = self.get_rank(node)
            if rank != INF:
                item = HeapItem(rank, node)
                heapq.heappush(heap, item)
            node = node.nxt
        
        
        # x → a → b
        while heap:
            heap_item = heapq.heappop(heap)
            a = heap_item.node
            b = a.nxt
                
            if not a.value.alive or b is None or heap_item.rank != self.get_rank(heap_item.node):
                continue
            
            pair = (self.vocab[a.value.vocab_id], self.vocab[b.value.vocab_id])
            new_token = pair[0] + pair[1]
            new_vocab_id = self.rev_vocab[new_token]
            
            # new_token 替代 (a, b)
            a.value = LinkItem(new_vocab_id)
            b.alive = False
            linkedlist.delete_node(b)
            
            
            x = a.pre
            if x is not None:
                if (rank := self.get_rank(x)) != INF:
                    heapq.heappush(heap, HeapItem(rank, x))
            
            if (rank := self.get_rank(a)) != INF:
                heapq.heappush(heap, HeapItem(rank, a))
            
        out = []
        node = linkedlist.head
        while node is not None:
            out.append(node.value.vocab_id)
            node = node.nxt
        return out
        
    
    def get_rank(self, node: LinkNode[LinkItem]) -> int:
        if not node.value.alive or node.nxt is None:
            return INF
        return self.merges_rank.get((self.vocab[node.value.vocab_id], self.vocab[node.nxt.value.vocab_id]), INF)

class Linear(nn.Module):
    
    def __init__(
        self, 
        in_features: int,
        out_features: int,
        device: Optional[torch.device]=None,
        dtype: Optional[torch.dtype]=None
    ):
        """
        计算时候是 y=xW^T,存储的是 W 而不是 W^T,所以要反过来
        """
        super(Linear, self).__init__()
        factory_kwargs = {"device": device, "dtype": dtype}
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.rand((out_features, in_features), **factory_kwargs))
        
        self.init_parameter()
    
    
    def init_parameter(self):
        std = math.sqrt(2 / (self.in_features + self.out_features))
        nn.init.trunc_normal_(self.weight, mean=0, std=std, a=-3*std, b=3*std)
    
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x 这里是行向量，所以进行的变化是 y=xW^T
        """
        return torch.matmul(x, self.weight.T)

class SwiGLU_FFN(nn.Module):
    
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        device: Optional[torch.device]=None,
        dtype: Optional[torch.dtype]=None
    ):
        super(SwiGLU_FFN, self).__init__()
        factory_kwargs = {"device": device, "dtype": dtype}
      
        self.d_ff = d_ff
        self.d_model = d_model
        
        self.w1 = nn.Parameter(torch.rand((d_ff, d_model), **factory_kwargs))
        self.w2 = nn.Parameter(torch.rand((d_model, d_ff), **factory_kwargs))
        self.w3 = nn.Parameter(torch.rand((d_ff, d_model), **factory_kwargs))
        
        self.init_parameter()
        
    def init_parameter(self):
        nn.init.normal_(self.w1, mean=0, std=math.sqrt(2/self.d_ff))
        nn.init.normal_(self.w2, mean=0, std=math.sqrt(2/self.d_model))
        nn.init.normal_(self.w3, mean=0, std=math.sqrt(2/self.d_ff))
    
    
    def SiLU(self, x: torch.Tensor):
        return x * torch.sigmoid(x)
    
    
    def forward(self, x: torch.Tensor):
        return (self.SiLU(x @ self.w1.T) * (x @ self.w3.T)) @ self.w2.T
        

class RoPE(nn.Module):
    
    def __init__(
        self,
        theta: float,
        d_k: int,
        max_seq_len: int,
        device: Optional[torch.device]=None
    ):
        super(RoPE, self).__init__()
        
        position = torch.arange(max_seq_len, device=device, dtype=torch.float32) # [max_len]
        freq = 1.0 / (theta ** (torch.arange(0, d_k, 2, device=device, dtype=torch.float32) / d_k))
        sinusoid = torch.outer(position, freq)
        rot = torch.exp(1j * sinusoid)  # [max_seq_len, dim//2]
        self.register_buffer("rot_cache", rot)
    
    def forward(self, qk: torch.Tensor, token_positions: torch.Tensor):
        *_, dim = qk.shape
    
        assert dim % 2 == 0, "dim must be even"
        
        qk_complex = qk.view(*qk.shape[:-1], dim//2, 2) # [bsize, nheads, seq_len, dim//2, 2]
        qk_complex = torch.view_as_complex(qk_complex)  # [bsize, nheads, seq_len, dim//2]
        
       
        rotated_qk_complex = qk_complex * self.rot_cache[token_positions]
        rotated_qk = torch.view_as_real(rotated_qk_complex)  # [bsize, nheads, seq_len, dim//2, 2]
        rotated_qk = rotated_qk.view_as(qk)
        
        return rotated_qk

class Multihead_Self_Attention(nn.Module):

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        use_rope: bool = False,
        device: Optional[torch.device] = None,
        **rope_config
    ):
        super(Multihead_Self_Attention, self).__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.use_rope = use_rope

        if use_rope:
            self.rope = RoPE(**rope_config, device=device)

        self.q_weight = nn.Parameter(torch.rand((d_model, d_model), device=device))
        self.k_weight = nn.Parameter(torch.rand((d_model, d_model), device=device))
        self.v_weight = nn.Parameter(torch.rand((d_model, d_model), device=device))
        self.o_weight = nn.Parameter(torch.rand((d_model, d_model), device=device))

    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor],
        token_positions: Optional[torch.Tensor] = None,
    ):

        q = x @ self.q_weight.T
        k = x @ self.k_weight.T
        v = x @ self.v_weight.T

        q = q.view(*x.shape[:-1], self.num_heads, self.d_k).transpose(1, 2)
        k = k.view(*x.shape[:-1], self.num_heads, self.d_k).transpose(1, 2)
        v = v.view(*x.shape[:-1], self.num_heads, self.d_k).transpose(1, 2)

        if self.use_rope:
            q = self.rope(q, token_positions)
            k = self.rope(k, token_positions)

        attn = scaled_dot_production_attention(q, k, v, mask)
        attn = attn.transpose(1, 2).contiguous().view_as(x)
        return attn @ self.o_weight.T

def run_transformer_block(
    d_model: int,
    num_heads: int,
    d_ff: int,
    max_seq_len: int,
    theta: float,
    weights: dict[str, Tensor],
    in_features: Float[Tensor, " batch sequence_length d_model"],
) -> Float[Tensor, " batch sequence_length d_model"]:
    from infra.transformer import TransformerBlock
    assert d_model % num_heads == 0

    device, dtype = in_features.device, in_features.dtype
    d_k = d_model // num_heads
    block = TransformerBlock(
        d_model,
        d_ff,
        num_heads, 
        use_rope=True, 
        device=device,
        dtype=dtype,
        theta=theta,
        max_seq_len=max_seq_len,
        d_k=d_k,
    )
    weights["attn.o_proj.weight"] = weights["attn.output_proj.weight"]
    block.load_state_dict(weights, strict=False)
    
    mask = torch.tril(torch.ones((max_seq_len, max_seq_len), device=device))
    B, S, _ = in_features.shape
    token_positions = torch.arange(S, device=device).expand(B, -1)
    return block(in_features, mask, token_positions)

class Transformer(nn.Module):
    
    def __init__(
        self,
        vocab_size: int,
        context_length: int,
        d_model: int,
        d_ff: int,
        num_layers: int,
        num_heads: int,
        device: Optional[torch.device],
        dtype: Optional[torch.dtype],
        theta: float,
        d_k: int,
        eps: float = 1e-5,
        use_rope: bool = True,
    ):
        super(Transformer, self).__init__()

        self.device = device
        self.context_length = context_length
        
        
        self.token_embeddings = Embedding(vocab_size, d_model, device, dtype)
        self.layers = nn.ModuleList([
                TransformerBlock(d_model, d_ff, num_heads, use_rope, device, dtype, theta, d_k, context_length, eps)
            for _ in range(num_layers)])
        self.ln_final = RMSNorm(d_model, eps, device, dtype)
        self.lm_head = Linear(d_model, vocab_size)
        
    def forward(self, x: torch.Tensor):
        B, L = x.shape
        positions = torch.arange(L, device=self.device).unsqueeze(0).expand(B, L)
        
        embed = self.token_embeddings(x)
        attn = embed
        for layer in self.layers:
            attn = layer(attn, positions)
        
        norm = self.ln_final(attn)
        proj = self.lm_head(norm)
        
        # return Softmax(proj, dim=-1)
        return proj