Assignment1

def run_train_bpe(
    input_path: str | os.PathLike,
    vocab_size: int,
    special_tokens: list[str],
    **kwargs,
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
    """Given the path to an input corpus, run train a BPE tokenizer and
    output its vocabulary and merges.

    Args:
        input_path (str | os.PathLike): Path to BPE tokenizer training data.
        vocab_size (int): Total number of items in the tokenizer's vocabulary (including special tokens).
        special_tokens (list[str]): A list of string special tokens to be added to the tokenizer vocabulary.
            These strings will never be split into multiple tokens, and will always be
            kept as a single token. If these special tokens occur in the `input_path`,
            they are treated as any other string.

    Returns:
        tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
            vocab:
                The trained tokenizer vocabulary, a mapping from int (token ID in the vocabulary)
                to bytes (token bytes)
            merges:
                BPE merges. Each list item is a tuple of bytes (<token1>, <token2>),
                representing that <token1> was merged with <token2>.
                Merges are ordered by order of creation.
    """

    # validation
    if not isinstance(vocab_size, int) or vocab_size <= 0:
        raise ValueError("vocab_size must be a positive integer")

    # -------------------------
    # 1. Initialize vocab
    # -------------------------
    vocab: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
    merges: list[tuple[bytes, bytes]] = []

    # add special tokens to vocab
    cnt_id = 256
    token_set = set(vocab.values())
    for s in special_tokens:
        if len(vocab) >= vocab_size: break
        s2bytes = s.encode("utf-8")
        if s2bytes not in token_set:
            token_set.add(s2bytes)
            vocab[cnt_id] = s2bytes
            cnt_id += 1

    if len(vocab) >= vocab_size:
        return vocab, merges

    # -------------------------
    # 2. Load corpus
    # -------------------------
    try:
        with open(input_path, "r", encoding="utf-8", errors="ignore") as f:
            text = f.read() # the whole file
    except FileNotFoundError:
        text = ""

    # -------------------------
    # 3. Pretokenization
    # -------------------------
    chunks = regex.split('|'.join(map(regex.escape, special_tokens)), text)
    # re
    PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
    token_freq = Counter()
    for chunk in chunks:
        for word in regex.findall(PAT, chunk):
            word_bytes = word.encode("utf-8")
            bytes_lst = [bytes([x]) for x in word_bytes]  #e.g. ['h', 'e', 'l', 'l', 'o']
            token_freq[tuple(bytes_lst)] += 1

    # -------------------------
    # 4. Initialize pair_freq
    # -------------------------
    pair_freq: Counter[tuple[bytes, bytes]] = Counter()
    for token, freq in token_freq.items():
        for l, r in pairwise(token):
            pair_freq[(l, r)] += freq

    # -------------------------
    # 5. Max heap with lazy deletion
    # -------------------------
    class MaxHeapItem:
        def __init__(self, cnt: int, p1: bytes, p2: bytes):
            self.cnt = cnt
            self.p1 = p1
            self.p2 = p2

        def __lt__(self, other):
            # Python heapq 是小根堆，这里反过来实现大根堆
            # tie-breaking: larger pair wins lexicographically
            return (self.cnt, self.p1, self.p2) > (other.cnt, other.p1, other.p2)

        def __eq__(self, other):
            return (self.cnt, self.p1, self.p2) == (other.cnt, other.p1, other.p2)

        @property
        def pair(self) -> tuple[bytes, bytes]:
            return self.p1, self.p2

        def __repr__(self):
            return f"({self.cnt}, ({self.p1}, {self.p2}))"

    # (cnt, p1, p2)
    lzheap: list[MaxHeapItem] = []
    for (l, r), c in pair_freq.items():
        heapq.heappush(lzheap, MaxHeapItem(c, l, r))

    # -------------------------
    # Helper functions
    # -------------------------
    def contains_pair(
        seq: tuple[bytes, ...],
        target_pair: tuple[bytes, bytes],
    ) -> bool:
        for i in range(len(seq) - 1):
            if seq[i] == target_pair[0] and seq[i + 1] == target_pair[1]:
                return True
        return False

    # new seq after merge
    def merge_seq_once(
        seq: tuple[bytes, ...],
        target_pair: tuple[bytes, bytes],
    ) -> tuple[bytes, ...]:
        """
        Left-to-right non-overlapping merge.

        Example:
            seq = (A, A, A), target = (A, A)
            result = (AA, A), not (AA, AA)
        """
        merged = []
        i = 0

        while i < len(seq):
            if i < len(seq) - 1 \
                and seq[i] == target_pair[0] \
                and seq[i + 1] == target_pair[1]:
                merged.append(seq[i] + seq[i + 1])
                i += 2
            else:
                merged.append(seq[i])
                i += 1

        return tuple(merged)

    # upd pair Counter
    def add_pair_count(pair: tuple[bytes, bytes], amount: int):
        pair_freq[pair] += amount
    def sub_pair_count(pair: tuple[bytes, bytes], amount: int):
        new_count = pair_freq[pair] - amount
        if new_count <= 0:
            del pair_freq[pair]
        else:
            pair_freq[pair] = new_count

    # -------------------------
    # 6. Main BPE loop
    # -------------------------
    while len(vocab) < vocab_size:
        best_item = None

        while lzheap:
            item = heapq.heappop(lzheap)
            pair = item.pair
            # check
            if pair in pair_freq and pair_freq[pair] == item.cnt:
                best_item = item
                break

        if best_item is None:
            break

        cur_pair = best_item.pair

        # record merge
        merges.append(cur_pair)
        # add new mapping
        new_token_bytes = cur_pair[0] + cur_pair[1]
        vocab[cnt_id] = new_token_bytes
        cnt_id += 1

        # all seqs containing cur_pair
        affected_seqs = [
            seq for seq in token_freq.keys()
            if contains_pair(seq, cur_pair)
        ]

        updated_pairs: set[tuple[bytes, bytes]] = set()
        new_seq_freq_delta: Counter[tuple[bytes, ...]] = Counter()

        for old_seq in affected_seqs:
            freq = token_freq[old_seq]

            # remove old sequence pair contributions
            for old_pair in pairwise(old_seq):
                updated_pairs.add(old_pair)
                sub_pair_count(old_pair, freq)

            # remove old sequence from token_freq
            del token_freq[old_seq]

            # merge old sequence
            new_seq = merge_seq_once(old_seq, cur_pair)

            # defer adding new sequence to token_freq
            new_seq_freq_delta[new_seq] += freq

        # add all new token sequences
        for new_seq, freq in new_seq_freq_delta.items():
            token_freq[new_seq] += freq

            # add new sequence pair contributions
            for new_pair in pairwise(new_seq):
                updated_pairs.add(new_pair)
                add_pair_count(new_pair, freq)

        # Push updated true counts into heap
        for pair in updated_pairs:
            if pair in pair_freq:
                heapq.heappush(
                    lzheap,
                    MaxHeapItem(pair_freq[pair], pair[0], pair[1]),
                )

    return (vocab, merges)

import json
import ast
import regex
from itertools import pairwise
from typing import Iterable, Iterator

PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""


class Tokenizer:
    def __init__(
        self,
        vocab: dict[int, bytes],
        merges: list[tuple[bytes, bytes]],
        special_tokens: list[str] = None,
    ):
        self.vocab = dict(vocab)
        self.merges = merges
        self.special_tokens = special_tokens or []
        self.cache = {}

        token_set = set(self.vocab.values())
        next_id = max(self.vocab.keys()) + 1 if self.vocab else 0

        # append user-provided special tokens if they are not already in vocab
        for token in self.special_tokens:
            token_bytes = token.encode("utf-8")
            if token_bytes not in token_set:
                self.vocab[next_id] = token_bytes
                token_set.add(token_bytes)
                next_id += 1

        self.bytes2id = {bts: id for id, bts in self.vocab.items()}
        self.merge_priority = {merge: i for i, merge in enumerate(self.merges)}

    @classmethod
    def from_files(
        cls,
        vocab_filepath: str,
        merges_filepath: str,
        special_tokens: list[str] = None,
    ):
        def bytes_to_unicode():
            bs = (
                list(range(ord("!"), ord("~") + 1))
                + list(range(ord("¡"), ord("¬") + 1))
                + list(range(ord("®"), ord("ÿ") + 1))
            )
            cs = bs[:]
            n = 0

            for b in range(256):
                if b not in bs:
                    bs.append(b)
                    cs.append(256 + n)
                    n += 1

            cs = [chr(c) for c in cs]
            return dict(zip(bs, cs))

        byte_encoder = bytes_to_unicode()
        byte_decoder = {v: k for k, v in byte_encoder.items()}

        def gpt2_str_to_bytes(s: str):
            return bytes(byte_decoder[c] for c in s)

        def to_bytes(x):
            if isinstance(x, bytes):
                return x

            if isinstance(x, list):
                return bytes(x)

            if isinstance(x, int):
                return bytes([x])

            if isinstance(x, str):
                if x.startswith("b'") or x.startswith('b"'):
                    try:
                        y = ast.literal_eval(x)
                        if isinstance(y, bytes):
                            return y
                    except Exception:
                        pass

                try:
                    return gpt2_str_to_bytes(x)
                except Exception:
                    pass

                try:
                    return x.encode("latin-1")
                except UnicodeEncodeError:
                    return x.encode("utf-8")

            raise ValueError(f"Cannot convert {x} to bytes")

        # load vocab
        with open(vocab_filepath, "r", encoding="utf-8") as f:
            raw_vocab = json.load(f)

        vocab: dict[int, bytes] = {}

        if isinstance(raw_vocab, dict):
            for k, v in raw_vocab.items():
                # format: {"0": [0], "1": [1], ...}
                if isinstance(k, str) and k.isdigit():
                    vocab[int(k)] = to_bytes(v)

                # GPT-2 format: {"token": id}
                elif isinstance(v, int):
                    vocab[v] = to_bytes(k)

                else:
                    raise ValueError("Unsupported vocab format")

        elif isinstance(raw_vocab, list):
            # format: [[0, [0]], [1, [1]], ...]
            for item in raw_vocab:
                if isinstance(item, list) and len(item) == 2:
                    idx, token = item
                    vocab[int(idx)] = to_bytes(token)
                else:
                    raise ValueError("Unsupported vocab format")

        else:
            raise ValueError("Unsupported vocab format")

        # load merges
        merges: list[tuple[bytes, bytes]] = []

        with open(merges_filepath, "r", encoding="utf-8") as f:
            lines = f.readlines()

        for line in lines:
            line = line.strip()

            if not line:
                continue

            # GPT-2 merges.txt 第一行通常是版本声明
            if line.startswith("#"):
                continue

            try:
                item = ast.literal_eval(line)
                if len(item) != 2:
                    raise ValueError("Unsupported merges format")
                merges.append((to_bytes(item[0]), to_bytes(item[1])))
                continue
            except Exception:
                pass

            parts = line.split()

            if len(parts) != 2:
                raise ValueError("Unsupported merges format")

            merges.append((to_bytes(parts[0]), to_bytes(parts[1])))

        return cls(vocab, merges, special_tokens)

    def _bpe_merge(self, words: bytes) -> list[bytes]:
        if words in self.cache:
            return self.cache[words]

        wordsbytes = [bytes([x]) for x in words]
        merge_priority = self.merge_priority

        while len(wordsbytes) > 1:
            good_pairs = set(
                (l, r) for l, r in pairwise(wordsbytes)
                if (l, r) in merge_priority
            )

            if not good_pairs:
                break

            best_pair = min(good_pairs, key=lambda x: merge_priority[x])

            # O(1) space implementation
            i = 0
            for x in wordsbytes:
                wordsbytes[i] = x
                i += 1

                if i > 1 \
                    and wordsbytes[i - 2] == best_pair[0] \
                    and wordsbytes[i - 1] == best_pair[1]:
                    wordsbytes[i - 2] += wordsbytes[i - 1]
                    i -= 1

            del wordsbytes[i:]

        self.cache[words] = wordsbytes
        return wordsbytes

    def encode(self, text: str) -> list[int]:
        if not text:
            return []

        special_tokens = self.special_tokens
        bytes2id = self.bytes2id

        if special_tokens:
            special_tokens_sorted = sorted(special_tokens, key=len, reverse=True)
            special_pattern = "|".join(map(regex.escape, special_tokens_sorted))
            chunks = regex.split(f"({special_pattern})", text)
        else:
            chunks = [text]

        ids = []

        for chunk in chunks:
            if not chunk:
                continue

            if chunk in special_tokens:
                ids.append(bytes2id[chunk.encode("utf-8")])
                continue

            for word in regex.findall(PAT, chunk):
                if not word:
                    continue

                merged_word = self._bpe_merge(word.encode("utf-8"))

                for s in merged_word:
                    ids.append(bytes2id[s])

        return ids

    def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]:
        for text in iterable:
            yield from self.encode(text)

    def decode(self, ids: list[int]) -> str:
        all2bytes = b"".join(self.vocab[id] for id in ids)
        return all2bytes.decode("utf-8", errors="replace")

import torch
from torch import nn

class Linear(nn.Module):
    def __init__(self, in_features: int, out_features: int, device: torch.device=None, dtype: torch.dtype=None):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.device = device
        self.dtype = dtype

        self.W = nn.Parameter(torch.rand(out_features, in_features, device=device, dtype=dtype))
        # no bias
        # approximate initializations given in handout
        std = (2 / (in_features + out_features)) ** 0.5
        nn.init.trunc_normal_(self.W, std=std, a=-3*std, b=3*std)

    def forward(self, x: torch.tensor) -> torch.tensor:
        return x @ self.W.T
    

def run_linear(
    d_in: int,
    d_out: int,
    weights: Float[Tensor, " d_out d_in"],
    in_features: Float[Tensor, " ... d_in"],
) -> Float[Tensor, " ... d_out"]:
    """
    Given the weights of a Linear layer, compute the transformation of a batched input.

    Args:
        in_dim (int): The size of the input dimension
        out_dim (int): The size of the output dimension
        weights (Float[Tensor, "d_out d_in"]): The linear weights to use
        in_features (Float[Tensor, "... d_in"]): The output tensor to apply the function to

    Returns:
        Float[Tensor, "... d_out"]: The transformed output of your linear module.
    """
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    mlp = Linear(in_features=d_in, out_features=d_out, device=device, dtype=torch.float32)
    mlp.W = nn.Parameter(weights)
    return mlp(in_features)

import torch
from torch import nn

class Embedding(nn.Module):
    def __init__(self, num_embeddings: int, embedding_dim: int, device: torch.device=None, dtype: torch.dtype=None):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.device = device
        self.dtype = dtype

        # also row main vector
        self.embeddings = nn.Parameter(torch.rand(num_embeddings, embedding_dim, device=device, dtype=dtype))
        std = 1
        nn.init.trunc_normal_(self.embeddings, std=std, a=-3*std, b=3*std)


    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        return self.embeddings[token_ids]
    

def run_embedding(
    vocab_size: int,
    d_model: int,
    weights: Float[Tensor, " vocab_size d_model"],
    token_ids: Int[Tensor, " ..."],
) -> Float[Tensor, " ... d_model"]:
    """
    Given the weights of an Embedding layer, get the embeddings for a batch of token ids.

    Args:
        vocab_size (int): The number of embeddings in the vocabulary
        d_model (int): The size of the embedding dimension
        weights (Float[Tensor, "vocab_size d_model"]): The embedding vectors to fetch from
        token_ids (Int[Tensor, "..."]): The set of token ids to fetch from the Embedding layer

    Returns:
        Float[Tensor, "... d_model"]: Batch of embeddings returned by your Embedding layer.
    """
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    emb = Embedding(vocab_size, d_model, device=device, dtype=torch.float32)
    emb.embeddings = nn.Parameter(weights)
    return emb(token_ids)

import torch
from torch import nn

class RMSNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1E-5, device: torch.device = None, dtype: torch.dtype = None):
        super().__init__()
        self.d_model = d_model
        self.eps = eps
        self.device = device
        self.dtype = dtype
        self.g = nn.Parameter(torch.ones(d_model, device=device, dtype=dtype))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_type = x.dtype
        x = x.to(torch.float32)
        den = (x**2).mean(dim=-1, keepdim=True)
        x = x / torch.sqrt(den + self.eps)
        # 注意是hadamard乘积
        return (self.g * x).to(x_type)

def run_rmsnorm(
    d_model: int,
    eps: float,
    weights: Float[Tensor, " d_model"],
    in_features: Float[Tensor, " ... d_model"],
) -> Float[Tensor, " ... d_model"]:
    """Given the weights of a RMSNorm affine transform,
    return the output of running RMSNorm on the input features.

    Args:
        d_model (int): The dimensionality of the RMSNorm input.
        eps: (float): A value added to the denominator for numerical stability.
        weights (Float[Tensor, "d_model"]): RMSNorm weights.
        in_features (Float[Tensor, "... d_model"]): Input features to run RMSNorm on. Can have arbitrary leading
            dimensions.

    Returns:
        Float[Tensor,"... d_model"]: Tensor of with the same shape as `in_features` with the output of running
        RMSNorm of the `in_features`.
    """
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    norm = RMSNorm(d_model=d_model, eps=eps, device=device, dtype=torch.float32)
    norm.g = nn.Parameter(weights)
    return norm(in_features)

import torch
from torch import nn
from Linear import Linear

# 注意运算都是 Hadamard 乘积
class SwiGLUFFN(nn.Module):
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.d_model = d_model
        self.d_ff = d_ff
        self.W1 = Linear(d_model, d_ff)
        self.W2 = Linear(d_ff, d_model)
        self.W3 = Linear(d_model, d_ff)
        
    def forward(self, x: torch.Tensor):
        w1x = self.W1(x)
        w3x = self.W3(x)
        return self.W2(self._SiLU(w1x) * w3x)

    def _SiLU(self, x: torch.Tensor):
        return x * torch.sigmoid(x)

def run_swiglu(
    d_model: int,
    d_ff: int,
    w1_weight: Float[Tensor, " d_ff d_model"],
    w2_weight: Float[Tensor, " d_model d_ff"],
    w3_weight: Float[Tensor, " d_ff d_model"],
    in_features: Float[Tensor, " ... d_model"],
) -> Float[Tensor, " ... d_model"]:
    """Given the weights of a SwiGLU network, return
    the output of your implementation with these weights.

    Args:
        d_model (int): Dimensionality of the feedforward input and output.
        d_ff (int): Dimensionality of the up-project happening internally to your swiglu.
        w1_weight (Float[Tensor, "d_ff d_model"]): Stored weights for W1
        w2_weight (Float[Tensor, "d_model d_ff"]): Stored weights for W2
        w3_weight (Float[Tensor, "d_ff d_model"]): Stored weights for W3
        in_features (Float[Tensor, "... d_model"]): Input embeddings to the feed-forward layer.

    Returns:
        Float[Tensor, "... d_model"]: Output embeddings of the same shape as the input embeddings.
    """
    # Example:
    # If your state dict keys match, you can use `load_state_dict()`
    # swiglu.load_state_dict(weights)
    # You can also manually assign the weights
    # swiglu.w1.weight.data = w1_weight
    # swiglu.w2.weight.data = w2_weight
    # swiglu.w3.weight.data = w3_weight
    ffn = SwiGLUFFN(d_model=d_model, d_ff=d_ff)
    ffn.W1.W = nn.Parameter(w1_weight)
    ffn.W2.W = nn.Parameter(w2_weight)
    ffn.W3.W = nn.Parameter(w3_weight)
    return ffn(in_features)

self.register_buffer(persistent=False)

nn.Parameter

import torch
from torch import nn

class RoPE(nn.Module):
    def __init__(self, theta: float, d_k: int, max_seq_len: int, device: torch.device = None):
        super().__init__()

        self.theta = theta

        if d_k % 2 > 0:
            raise ValueError("d_k must be an even positive integer")

        self.theta = theta
        self.d_k = d_k  # dimension of query and key
        self.max_seq_len = max_seq_len
        self.device = device

        exps = 1 / (theta ** (torch.arange(0, d_k, 2, dtype=torch.float32, device=device) / d_k))
        positions = torch.arange(max_seq_len, device=device)
        # outer(a[0..n), b[0..n)) =>c_ij = ai * bj
        vals = torch.outer(positions, exps)
        self.register_buffer("cos_memo", vals.cos(), persistent=False)
        self.register_buffer("sin_memo", vals.sin(), persistent=False)

    def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor:
        x = x.to(self.device)

        x_e = x[..., 0::2]
        x_o = x[..., 1::2]

        cos = self.cos_memo[token_positions]
        sin = self.sin_memo[token_positions]

        # [1, max_seq_len, d_k//2] => [max_seq_len, d_k//2]
        # to make use of broadcast
        cos = cos.unsqueeze(0)
        sin = sin.unsqueeze(0)    

        # [batch, max_seq_len, d_k//2]
        new_x_e = x_e * cos - x_o * sin
        new_x_o = x_o * cos + x_e * sin

        # [batch, max_seq_len, d_k//2,2]
        res = torch.stack([new_x_e, new_x_o], dim=-1)
    
        # [batch, max_seq_len, d_k], 然后flatten恰好会把奇偶交错排列
        return res.flatten(-2)
    

def run_rope(
    d_k: int,
    theta: float,
    max_seq_len: int,
    in_query_or_key: Float[Tensor, " ... sequence_length d_k"],
    token_positions: Int[Tensor, " ... sequence_length"],
) -> Float[Tensor, " ... sequence_length d_k"]:
    """
    Run RoPE for a given input tensor.

    Args:
        d_k (int): Embedding dimension size for the query or key tensor.
        theta (float): RoPE parameter.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        in_query_or_key (Float[Tensor, "... sequence_length d_k"]): Input tensor to run RoPE on.
        token_positions (Int[Tensor, "... sequence_length"]): Tensor of shape (batch_size, sequence_length) with the token positions
    Returns:
        Float[Tensor, " ... sequence_length d_k"]: Tensor with RoPEd input.
    """
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    rope = RoPE(theta=theta, d_k=d_k, max_seq_len=max_seq_len, device=device)
    return rope(in_query_or_key, token_positions)

import torch
from torch import nn

def softmax(x: torch.Tensor, dim: int) -> torch.Tensor:
    '''
    x: [batch_size, seq_len, d_model]
    x.max is a tuple: (values, indices)
    so x.max()[0] is a scalar, or 1x1 tensor
    '''

    ma = x.max(dim=dim, keepdim=True)[0]
    x_exp = torch.exp(x - ma)
    Z_theta = x_exp.sum(dim=dim, keepdim=True)
    return x_exp / Z_theta

def run_softmax(in_features: Float[Tensor, " ..."], dim: int) -> Float[Tensor, " ..."]:
    """
    Given a tensor of inputs, return the output of softmaxing the given `dim`
    of the input.

    Args:
        in_features (Float[Tensor, "..."]): Input features to softmax. Shape is arbitrary.
        dim (int): Dimension of the `in_features` to apply softmax to.

    Returns:
        Float[Tensor, "..."]: Tensor of with the same shape as `in_features` with the output of
        softmax normalizing the specified `dim`.
    """
    return softmax(in_features, dim)

[[True, True, False]]

import torch
from torch import nn
from softmax import softmax

def scaled_dot_product_attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, mask: torch.Tensor | None = None):
    d_k = Q.shape[-1]
    # batch mat mul and batch mat transpose
    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1E9)
    attn_weights = softmax(scores, dim=-1)
    return torch.matmul(attn_weights, V)

def run_scaled_dot_product_attention(
    Q: Float[Tensor, " ... queries d_k"],
    K: Float[Tensor, " ... keys d_k"],
    V: Float[Tensor, " ... keys d_v"],
    mask: Bool[Tensor, " ... queries keys"] | None = None,
) -> Float[Tensor, " ... queries d_v"]:
    """
    Given key (K), query (Q), and value (V) tensors, return
    the output of your scaled dot product attention implementation.

    Args:
        Q (Float[Tensor, " ... queries d_k"]): Query tensor
        K (Float[Tensor, " ... keys d_k"]): Key tensor
        V (Float[Tensor, " ... keys d_v"]): Values tensor
        mask (Bool[Tensor, " ... queries keys"] | None): Mask tensor
    Returns:
        Float[Tensor, " ... queries d_v"]: Output of SDPA
    """
    return scaled_dot_product_attention(Q, K, V, mask)

import torch
from torch import nn
from softmax import softmax

class multihead_self_attention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        self.d_model = d_model  # input dim
        self.num_heads = num_heads  # num of head
        self.head_dim = d_model // num_heads    # dim of each head
    
    def _attention(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, mask: torch.Tensor | None = None):
        d_k = self.head_dim # 注意，每个头的维度是 head_dim
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))
        if mask is not None:
            scores = scores.masked_fill(mask == 1, -1E9)
        attn_weights = softmax(scores, dim=-1)
        return torch.matmul(attn_weights, V)

    def forward(self, x: torch.Tensor, Wq: torch.Tensor, Wk: torch.Tensor, Wv: torch.Tensor, Wo: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, self.d_model = x.shape

        q = x @ Wq.T  # (batch_size, seq_len, d_model) @ (d_model, d_k) -> (batch_size, seq_len, d_k)
        k = x @ Wk.T  # (batch_size, seq_len, d_model) @ (d_model, d_k) -> (batch_size, seq_len, d_k)
        v = x @ Wv.T  # (batch_size, seq_len, d_model) @ (d_model, d_k) -> (batch_size, seq_len, d_k)
        
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim)
        
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # 对角线上方第一条对角往上都是1
        mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=x.device), diagonal=1)
        mask = mask.unsqueeze(0).unsqueeze(0) # [1, 1, seq_len, seq_len]

        attn = self._attention(q, k, v, mask) # [batch_size, num_heads, seq_len, self.head_dim]
        attn = attn.transpose(1, 2) # # [batch_size, seq_len, num_heads, self.head_dim]
        attn = attn.contiguous().view(batch_size, seq_len, self.d_model)
        return attn @ Wo.T

def run_multihead_self_attention(
    d_model: int,
    num_heads: int,
    q_proj_weight: Float[Tensor, " d_model d_model"],
    k_proj_weight: Float[Tensor, " d_model d_model"],
    v_proj_weight: Float[Tensor, " d_model d_model"],
    o_proj_weight: Float[Tensor, " d_model d_model"],
    in_features: Float[Tensor, " ... sequence_length d_model"],
) -> Float[Tensor, " ... sequence_length d_model"]:
    """
    Given the key, query, and value projection weights of a naive unbatched
    implementation of multi-head attention, return the output of an optimized batched
    implementation. This implementation should handle the key, query, and value projections
    for all heads in a single matrix multiply.
    This function should not use RoPE.
    See section 3.2.2 of Vaswani et al., 2017.

    Args:
        d_model (int): Dimensionality of the feedforward input and output.
        num_heads (int): Number of heads to use in multi-headed attention.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        q_proj_weight (Float[Tensor, "d_model d_model"]): Weights for the Q projection
        k_proj_weight (Float[Tensor, "d_model d_model"]): Weights for the K projection
        v_proj_weight (Float[Tensor, "d_model d_model"]): Weights for the V projection
        o_proj_weight (Float[Tensor, "d_model d_model"]): Weights for the output projection
        in_features (Float[Tensor, "... sequence_length d_model"]): Tensor to run your implementation on.

    Returns:
        Float[Tensor, " ... sequence_length d_model"]: Tensor with the output of running your optimized, batched multi-headed attention
        implementation with the given QKV projection weights and input features.
    """
    model = multihead_self_attention(d_model, num_heads)
    return model(in_features, q_proj_weight, k_proj_weight, v_proj_weight, o_proj_weight)

import torch
from torch import nn
from softmax import softmax
from RoPE import RoPE

class multihead_self_attention_rope(nn.Module):
    def __init__(self, d_model: int, num_heads: int, max_seq_len: int, theta: float, device: torch.device | None = None):
        super().__init__()
        self.d_model = d_model  # input dim
        self.num_heads = num_heads  # num of head
        self.head_dim = d_model // num_heads    # dim of each head
        self.rope = RoPE(theta, self.head_dim, max_seq_len, device)
    
    def _attention(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, mask: torch.Tensor | None = None):
        d_k = self.head_dim # 注意，每个头的维度是 head_dim
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
        if mask is not None:
            scores = scores.masked_fill(mask == 1, -1E9)
        attn_weights = softmax(scores, dim=-1)
        return torch.matmul(attn_weights, V)

    def forward(self, x: torch.Tensor, Wq: torch.Tensor, Wk: torch.Tensor, Wv: torch.Tensor, Wo: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, self.d_model = x.shape

        q = x @ Wq.T  # (batch_size, seq_len, d_model) @ (d_model, d_k) -> (batch_size, seq_len, d_k)
        k = x @ Wk.T  # (batch_size, seq_len, d_model) @ (d_model, d_k) -> (batch_size, seq_len, d_k)
        v = x @ Wv.T  # (batch_size, seq_len, d_model) @ (d_model, d_k) -> (batch_size, seq_len, d_k)
        
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim)
        
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        q = self.rope(q, token_positions)
        k = self.rope(k, token_positions)

        # 对角线上方第一条对角往上都是1
        mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=x.device), diagonal=1)
        mask = mask.unsqueeze(0).unsqueeze(0) # [1, 1, seq_len, seq_len]

        attn = self._attention(q, k, v, mask) # [batch_size, num_heads, seq_len, self.head_dim]
        attn = attn.transpose(1, 2) # # [batch_size, seq_len, num_heads, self.head_dim]
        attn = attn.contiguous().view(batch_size, seq_len, self.d_model)
        return attn @ Wo.T

def run_rope(
    d_k: int,
    theta: float,
    max_seq_len: int,
    in_query_or_key: Float[Tensor, " ... sequence_length d_k"],
    token_positions: Int[Tensor, " ... sequence_length"],
) -> Float[Tensor, " ... sequence_length d_k"]:
    """
    Run RoPE for a given input tensor.

    Args:
        d_k (int): Embedding dimension size for the query or key tensor.
        theta (float): RoPE parameter.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        in_query_or_key (Float[Tensor, "... sequence_length d_k"]): Input tensor to run RoPE on.
        token_positions (Int[Tensor, "... sequence_length"]): Tensor of shape (batch_size, sequence_length) with the token positions
    Returns:
        Float[Tensor, " ... sequence_length d_k"]: Tensor with RoPEd input.
    """
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    rope = RoPE(theta=theta, d_k=d_k, max_seq_len=max_seq_len, device=device)
    return rope(in_query_or_key, token_positions)

x = x / torch.sqrt(den + self.eps)

x /= torch.sqrt(den + self.eps)

import torch
from torch import nn
from SwiGLUFFN import SwiGLUFFN
from multihead_self_attention_rope import multihead_self_attention_rope
from RMSNorm import RMSNorm

class transformer_block(nn.Module):
    def __init__(self, d_model: int, num_heads: int, d_ff: int, max_seq_len: int, theta: float, attn_q_proj_weight: torch.Tensor, attn_k_proj_weight: torch.Tensor, attn_v_proj_weight: torch.Tensor, attn_o_proj_weight: torch.Tensor, ln1_weight: torch.Tensor, ln2_weight: torch.Tensor, ffn_w1_weight: torch.Tensor, ffn_w2_weight: torch.Tensor, ffn_w3_weight: torch.Tensor, device=None):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_ff = d_ff
        self.max_seq_len = max_seq_len
        self.device = device
        self.theta = theta

        self.attn_q_proj_weight = attn_q_proj_weight
        self.attn_k_proj_weight = attn_k_proj_weight
        self.attn_v_proj_weight = attn_v_proj_weight
        self.attn_o_proj_weight = attn_o_proj_weight

        self.rms1 = RMSNorm(d_model, eps=1E-5, device=device)
        self.rms2 = RMSNorm(d_model, eps=1E-5, device=device)
        self.rms1.load_state_dict({"g": ln1_weight})
        self.rms2.load_state_dict({"g": ln2_weight})

        self.ffn = SwiGLUFFN(d_model, d_ff)
        self.ffn.load_state_dict({"W1.W": ffn_w1_weight, "W2.W": ffn_w2_weight, "W3.W": ffn_w3_weight})

        self.MHA = multihead_self_attention_rope(d_model, num_heads, max_seq_len, theta, device)

    def forward(self, in_features: torch.Tensor) -> torch.Tensor:
        # [batch_size, seq_len, ...]
        # shape[1] = seq_len
        token_positions = torch.arange(in_features.shape[-2], device=in_features.device)
        x1 = self.rms1(in_features)
        x1 = self.MHA(x1, self.attn_q_proj_weight, self.attn_k_proj_weight, self.attn_v_proj_weight, self.attn_o_proj_weight, token_positions)

        x1 = in_features + x1

        x2 = self.rms2(x1)
        x2 = self.ffn(x2)
        return x1 + x2

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"]:
    """
    Given the weights of a pre-norm Transformer block and input features,
    return the output of running the Transformer block on the input features.

    This function should use RoPE.
    Depending on your implementation, you may simply need to pass the relevant args
    to your TransformerBlock constructor, or you may need to initialize your own RoPE
    class and pass that instead.

    Args:
        d_model (int): The dimensionality of the Transformer block input.
        num_heads (int): Number of heads to use in multi-headed attention. `d_model` must be
            evenly divisible by `num_heads`.
        d_ff (int): Dimensionality of the feed-forward inner layer.
        max_seq_len (int): Maximum sequence length to pre-cache if your implementation does that.
        theta (float): RoPE parameter.
        weights (dict[str, Tensor]):
            State dict of our reference implementation.
            The keys of this dictionary are:
            - `attn.q_proj.weight`
                The query projections for all `num_heads` attention heads.
                Shape is (d_model, d_model).
                The rows are ordered by matrices of shape (num_heads, d_k),
                so `attn.q_proj.weight == torch.cat([q_heads.0.weight, ..., q_heads.N.weight], dim=0)`.
            - `attn.k_proj.weight`
                The key projections for all `num_heads` attention heads.
                Shape is (d_model, d_model).
                The rows are ordered by matrices of shape (num_heads, d_k),
                so `attn.k_proj.weight == torch.cat([k_heads.0.weight, ..., k_heads.N.weight], dim=0)`.
            - `attn.v_proj.weight`
                The value projections for all `num_heads` attention heads.
                Shape is (d_model, d_model).
                The rows are ordered by matrices of shape (num_heads, d_v),
                so `attn.v_proj.weight == torch.cat([v_heads.0.weight, ..., v_heads.N.weight], dim=0)`.
            - `attn.output_proj.weight`
                Weight of the multi-head self-attention output projection
                Shape is (d_model, d_model).
            - `ln1.weight`
                Weights of affine transform for the first RMSNorm
                applied in the transformer block.
                Shape is (d_model,).
            - `ffn.w1.weight`
                Weight of the first linear transformation in the FFN.
                Shape is (d_ff, d_model).
            - `ffn.w2.weight`
                Weight of the second linear transformation in the FFN.
                Shape is (d_model, d_ff).
            - `ffn.w3.weight`
                Weight of the third linear transformation in the FFN.
                Shape is (d_ff, d_model).
            - `ln2.weight`
                Weights of affine transform for the second RMSNorm
                applied in the transformer block.
                Shape is (d_model,).
        in_features (Float[Tensor, "batch sequence_length d_model"]):
            Tensor to run your implementation on.

    Returns:
        Float[Tensor, "batch sequence_length d_model"] Tensor with the output of
        running the Transformer block on the input features while using RoPE.
    """
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    dtype = in_features.dtype

    attn_q_proj_weight = weights["attn.q_proj.weight"].to(device=device, dtype=dtype)
    attn_k_proj_weight = weights["attn.k_proj.weight"].to(device=device, dtype=dtype)
    attn_v_proj_weight = weights["attn.v_proj.weight"].to(device=device, dtype=dtype)
    attn_o_proj_weight = weights["attn.output_proj.weight"].to(device=device, dtype=dtype)

    ln1_weight = weights["ln1.weight"].to(device=device, dtype=dtype)
    ln2_weight = weights["ln2.weight"].to(device=device, dtype=dtype)

    ffn_w1_weight = weights["ffn.w1.weight"].to(device=device, dtype=dtype)
    ffn_w2_weight = weights["ffn.w2.weight"].to(device=device, dtype=dtype)
    ffn_w3_weight = weights["ffn.w3.weight"].to(device=device, dtype=dtype)

    model = transformer_block(d_model, num_heads, d_ff, max_seq_len, theta, attn_q_proj_weight, attn_k_proj_weight, attn_v_proj_weight, attn_o_proj_weight, ln1_weight, ln2_weight, ffn_w1_weight, ffn_w2_weight, ffn_w3_weight, device)
    return model(in_features)