CS336-Assignment1

word_cnt = {
    (b'\x61', b'\x62', b'\x63'): 100,    # "abc" 出现 100 次
    (b'\x64', b'\x65'): 50               # "de" 出现 50 次
}

pair_cnt = {
    (b'\x61', b'\x62'): 150,  # ("a"+"b") 字节对出现 150 次
    (b'\x62', b'\x63'): 100,  # ("b"+"c") 字节对出现 100 次
}

vocab = {
    0: b'\x00',              # 基础字节 0-255
    255: b'\xff',
    256: b'<|endoftext|>',   # 特殊 token
}

merges = [
    (b'\x61', b'\x62'),  # 第 1 次合并：a+b→ab
    (b'\x62', b'\x63'),  # 第 2 次合并：b+c→bc
]

def train_bpe(input_path,vocab_size,special_tokens):
    """
    Train a Byte Pair Encoding (BPE) tokenizer on the input text.

    This function implements the BPE algorithm to learn a vocabulary of a specified size
    by iteratively merging the most frequent pairs of tokens in the training data.

    Args:
        input_path (str): Path to the input text file to train on.
        vocab_size (int): The desired final vocabulary size, including special tokens.
        special_tokens (list): List of special tokens to be preserved during tokenization
            (e.g., ['<PAD>', '<UNK>', '<EOS>']).

    Returns:
        tuple: A tuple containing:
            - vocab (dict): A dictionary mapping token IDs to token strings. The first
              entries are basic vocabulary (bytes and special tokens), followed by
              merged tokens.
            - merges (list): A list of tuples representing the merge operations in order,
              where each tuple contains the pair of tokens that were merged.

    Process:
        1. Reads text from the input file
        2. Splits text by special tokens to preserve them
        3. Counts word frequencies in the corpus
        4. Counts pairs of adjacent tokens
        5. Iteratively merges the most frequent pair until vocab_size is reached
        6. Returns the final vocabulary and merge operations

    Note:
        - Uses parallel processing for large chunk counts (>= 4) to improve performance
        - The base vocabulary size is determined by basic tokens and special tokens
        - Number of merges = vocab_size - base_vocab_size
    """

    text = read_text(input_path)
    chunks = split_by_special(text,special_tokens)

    # Only parallelize if chunk count is big enough
    if len(chunks) < 4: word_dicts = list(map(count_word, chunks))
    else: word_dicts = process_map(count_word, chunks, chunksize=100)

    word_cnt = merge_dicts(word_dicts)
    pair_cnt = count_pair(word_cnt)

    vocab = get_basic_vocab(special_tokens)
    base_vocab_size = len(vocab)
    n_merges=vocab_size-base_vocab_size

    merges = []
    for i in range(n_merges):
        max_pair = get_max_pair(pair_cnt)
        vocab[base_vocab_size+i] = max_pair[0]+max_pair[1]
        merges.append(max_pair)
        word_cnt, pair_cnt = update_cnt_fast(word_cnt,pair_cnt,max_pair)
    return vocab, merges

# 词表：token ID → 字节对象
vocab = {
    0: b'\x00',              # 基础字节 0-255
    255: b'\xff',
    256: b'he',              # 合并后的 token
    258: b'll',              # 合并后的 token
    266: b'<|endoftext|>',   # 特殊 token
}

# 反向词表：字节对象 → token ID（加速查询）
vocab_to_id = {
    b'\x00': 0,
    b'\xff': 255,
    b'he': 256,
    b'll': 258,
    b'<|endoftext|>': 266,
}

# 合并规则：有序的合并对列表（记录训练顺序）
merges = [
    (b'h', b'e'),      # 第 1 次合并：token ID 256
    (b'e', b'l'),      # 第 2 次合并：token ID 257
    (b'l', b'l'),      # 第 3 次合并：token ID 258
    ...
]

# 合并集合：快速查询某个对是否能合并
merges_set = {(b'h', b'e'), (b'e', b'l'), (b'l', b'l'), ...}

# encode 编码过程
text = "hello"
↓
tokens = [256, 258, 111]  # token ID 列表
# decode 解码过程
ids = [256, 258, 111]
↓
text = "hello"  # 字符串

class Tokenizer:
    """
    A Byte Pair Encoding (BPE) tokenizer with support for special tokens.
    This tokenizer converts text into sequences of token IDs using a vocabulary and merge rules.
    It supports special tokens that should not be merged during encoding.
    Attributes:
        vocab (dict[int, bytes]): Mapping from token IDs to byte sequences.
        merges (set[tuple[bytes, bytes]]): Set of valid merge pairs for BPE.
        special_tokens (list[str]): List of special token strings (e.g., ["<|endoftext|>"]).
        special_tokens_bytes (list[bytes]): UTF-8 encoded versions of special tokens.
        vocab_to_id (dict[bytes, int]): Reverse mapping from byte sequences to token IDs.
    Args:
        vocab (dict[int, bytes]): Dictionary mapping token IDs to byte sequences.
        merges (list[tuple[bytes, bytes]]): List of byte pair merge rules.
        special_tokens (list[str], optional): List of special tokens to preserve during tokenization.
		"""
    def __init__(self, vocab, merges, special_tokens=None):
        self.vocab = vocab
        self.merges = set(merges)
        self.special_tokens = special_tokens if special_tokens else []
        self.special_tokens_bytes = [i.encode('utf-8') for i in self.special_tokens]
        

        self.vocab_to_id={v:k for k,v in vocab.items()}

        # Ensure special tokens are in the vocabulary
        for token_bytes in self.special_tokens_bytes:
            if token_bytes not in self.vocab_to_id:
                # Add to vocab if not already present
                new_id = len(self.vocab)
                self.vocab[new_id] = token_bytes
                self.vocab_to_id[token_bytes] = new_id

    @classmethod
    def from_files(cls, vocab_filepath, merges_filepath, special_tokens=None):
        """Loads from the clean Hex text format"""
        vocab = {}
        with open(f"{vocab_filepath}", 'r') as f:
            for line in f:
                idx, hx = line.strip().split('\t')
                vocab[int(idx)] = bytes.fromhex(hx)
        
        merges = []
        with open(f"{merges_filepath}", 'r') as f:
            for line in f:
                p1, p2 = line.strip().split('\t')
                merges.append((bytes.fromhex(p1), bytes.fromhex(p2)))
                
        return cls(vocab=vocab, merges=merges, special_tokens=special_tokens)

    def encode(self, text: str) -> list[int]:
        chunks = split_by_special(text, self.special_tokens, drop_special=False)
        tokens = []
        for chunk in chunks:
            if self.special_tokens and chunk in self.special_tokens:
                tokens.append(self.vocab_to_id[chunk.encode('utf-8')])
            else:
                tokens.extend(encode_merged(chunk, self.merges, self.vocab_to_id))
        return tokens

    def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]:
        """
        Given an iterable of strings (e.g., a Python file handle), return a generator that lazily yields token IDs. 
        This is required for memory-efficient tokenization of large files that we cannot directly load into memory.
        """
        for chunk in iterable:
            yield from self.encode(chunk)

    def decode(self, ids: list[int]) -> str:
        "Decode a sequence of token IDs into text."
        return b''.join([self.vocab[t] for t in ids]).decode('utf-8',errors='replace')
           
def apply_merges(word_bytes, merges_set, vocab_to_id):
  word_bytes = list(word_bytes)

  while True:
      min_token_id = float('inf')
      best_pair_idx = -1
      merged = None

      for i in range(len(word_bytes) - 1):
          pair = (word_bytes[i], word_bytes[i + 1])
          if pair in merges_set:
              combined = pair[0] + pair[1]
              token_id = vocab_to_id.get(combined)
              if token_id is not None and token_id < min_token_id:
                  min_token_id = token_id
                  best_pair_idx = i
                  merged = combined

      if best_pair_idx == -1:
          break

      # Apply best merge
      word_bytes = (
          word_bytes[:best_pair_idx]
          + [merged]
          + word_bytes[best_pair_idx + 2:]
      )

  return tuple(word_bytes)

def encode_merged(text,merges,vocab_to_id):
    word_list = split_to_words(text)
    tokens=[]
    for word in word_list:
        word_bytes=word2bytes(word)
        merged_word_bytes = apply_merges(word_bytes,merges,vocab_to_id)
        tokens.extend(vocab_to_id[i] for i in merged_word_bytes)
    return tokens

# 把角度变成 cos 和 sin
# 工业界有个 trick：为了方便后续计算，我们把 cos 和 sin 复制两份拼接起来
# 形状从 [end, dim/2] 变成 [end, dim],因为后续计算不会分成even和odd两部分，所以位置最大就是max_len
emb = torch.cat((freqs, freqs), dim=-1)

@staticmethod
def rotate_half(x):
    """
    这个辅助函数的作用是实现公式里的 (-y, x) 部分。
    假设输入 x 是 [x1, x2, ..., xd/2, y1, y2, ..., yd/2]
    """
    # 1. 把向量 x 最后一维切成两半
    x1 = x[..., : x.shape[-1] // 2] # 取前半部分 (相当于 x)
    x2 = x[..., x.shape[-1] // 2 :] # 取后半部分 (相当于 y)

    # 2. 拼接成 [-y, x]
    # 为什么是负号？因为旋转公式里有 (x*cos - y*sin)
    # 我们稍后会用 x * cos + (-y) * sin 来计算
    return torch.cat((-x2, x1), dim=-1)

# 2. 应用旋转公式（一行搞定！）
# 公式: x_rotated = x * cos + rotate_half(x) * sin
# 这等价于:
#   - 前半部分: x_even * cos - x_odd * sin
#   - 后半部分: x_even * sin + x_odd * cos
return x * cos + self.rotate_half(x) * sin

  def forward(
      self,
      x: Float[Tensor, "batch seq_len d_model"],
      token_positions: Int[Tensor, "batch seq_len"]| None = None,
  ) -> Float[Tensor, "batch seq_len d_model"]:
      B,S,D=x.shape
      q,k,v,o_proj=[nn.Linear(D,D) for _ in range(4)]
      q=q.view(B, S, self.num_heads, self.d_k).transpose(1, 2)
      k=k.view(B, S, self.num_heads, self.d_k).transpose(1, 2)
      v=v.view(B, S, self.num_heads, self.d_k).transpose(1, 2)
      if self.use_rope:
          q,k=self.rope(q,token_positions),self.rope(k,token_positions)
      attn_scores=q@k.transpose(-2, -1)/math.sqrt(self.d_k)
      if  self.causal_mask is not None:
          attn_scores=attn_scores.masked_fill(~self.causal_mask[:,:,:S,:S],float("-inf"))
      attn_probs=torch.softmax(attn_scores,dim=-1)
      out=attn_probs@v
      out= out.transpose(1, 2).contiguous().view(B, S, D)       
      return o_proj(out) 

def cross_entropy(logits,targets):
    vocab_size=logits.size(-1)
    logits=logits.view(-1,vocab_size)
    targets=targets.view(-1).to(torch.long)

    # 数值稳定性：减去每行最大值
    logits=logits-torch.max(logits,dim=-1,keepdim=True)[0]
    # log-sum-exp 技巧，shape:(N,vocab_size)-> (N,)
    log_sum_exp=torch.log(torch.sum(torch.exp(logits),dim=-1))
    # 选择目标类别的 logits
    # 这里使用二维的行索引与列索引，在二维矩阵中选取
    target_logits=logits[torch.arange(logits.size(0)),targets] 
    # 使用gather语法也可以，提取目标类别的 logit，但输出形状与 index 相同
    # 注意index shape应与logits一致二维，并且最终还要压缩还原为一维
    # target_logits=logits.gather(dim=-1,index=targets.unsqueeze(1)).squeeze(1)

    loss=log_sum_exp-target_logits # shape:(N,) - (N,)
    return loss.mean()

class AdamW(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9,0.999), weight_decay=1e-2, eps=1e-8):
        if lr < 0:
            raise ValueError(f"Invalid learning rate: {lr}")
        defaults = {"lr":lr, "beta": betas, "weight_decay": weight_decay, "eps": eps}
        super().__init__(params, defaults)
    def step(self):
        for group in self.param_groups:
            lr = group['lr']
            beta = group['beta']
            weight_decay = group['weight_decay']
            eps = group['eps']
            for p in group['params']:
                if p.grad is None:
                    continue
                state = self.state[p]
                m = state.get("m", torch.zeros_like(p.grad))
                v = state.get("v", torch.zeros_like(p.grad))
                m = beta[0] * m + (1 - beta[0]) * p.grad.data
                v = beta[1] * v + (1 - beta[1]) * p.grad.data ** 2
                t = state.get("t",1)
                lr_t = lr * ((1 - beta[1] ** t) ** 0.5) / (1 - beta[0] ** t)
                p.data -= lr_t * m / (v ** 0.5 + eps)
                p.data -= lr * weight_decay * p.data
                state["m"] = m
                state["v"] = v
                # 注意这里要自增1
                state["t"] = t + 1

    total_norm = 0
    for param in parameters:
        if param.grad is not None:
            total_norm += ((param.grad.data) ** 2).sum().item()
    total_norm = total_norm ** 0.5
    if total_norm > max_l2_norm:
        scale = max_l2_norm / (total_norm + 1e-6)
        for param in parameters:
            if param.grad is not None:
                param.grad.data = param.grad.data * scale
    return total_norm * scale if total_norm > max_l2_norm else total_norm

class BatchIterator:
    def __init__(self, data, batch_size, context_length, device="cpu"):
        """
        基于流式读取的 BatchIterator
        逻辑参考：nanoGPT / GPT-2 standard data loader
        """
        self.data = data # 预期是 mmap 的 numpy array
        self.B = batch_size
        self.T = context_length
        self.device = device
        self.current_position = 0
        self.total_len = len(data)
        
        print(f"BatchIterator initialized: {self.total_len:,} tokens.")

    def get_batch(self):
        B, T = self.B, self.T
        
        # 1. 边界检查：如果剩余数据不够一个 Batch，则重置 (Epoch 结束)
        # 我们需要 B*T + 1 个 token (因为要有 target)
        if self.current_position + B * T + 1 > self.total_len:
            self.current_position = 0
            print("Epoch finished, resetting to start of data.")
            
        # 2. 提取 Buffer
        # 从长流中直接切出一大块：大小为 B*T + 1
        buf = self.data[self.current_position : self.current_position + B * T + 1]
        
        # 3. 转换为 Tensor 并移至 GPU
        # 注意：先转 int64 确保兼容 Embedding 层
        buf = torch.from_numpy(buf.astype(np.int64)).to(self.device)
        
        # 4. 构建 Input (x) 和 Target (y)
        # buf[:-1] 是 0 到 BT-1
        # buf[1:]  是 1 到 BT
        # .view(B, T) 将一维流折叠成二维 Batch
        x = buf[:-1].view(B, T)
        y = buf[1:].view(B, T)
        
        # 5. 推进指针
        # 下一次读取从当前这一块的末尾开始，绝不重叠
        self.current_position += B * T
        
        return x, y

def get_batch(data,batch_size,context_length,device="cpu"):
    total_len=len(data)
    idx=np.random.randint(0,total_len-context_length,batch_size)

    inputs=np.array([data[i:i+context_length] for i in idx])
    targets=np.array([data[i+1:i+1+context_length] for i in idx])

    return (
        torch.from_numpy(inputs).to(device),
        torch.from_numpy(targets).to(device)
    )

def save_checkpoint(
        model:torch.nn.Module,
        optimizer:torch.optim.Optimizer,
        iteration:int,
        path:str
):
    save_dict={
        "model_state_dict":model.state_dict(),
        "optimzer_state_dict":optimizer.state_dict(),
        "iteration":iteration
    }
    torch.save(save_dict,path)

def load_checkpoint(
    src: str ,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer | None=None, # 推理时候不用optimizer
):
    state_dict=torch.load(src,map_location="cpu")
    
    model_state=state_dict["model_state_dict"]
    if any(k.startwith("_orig_mod.") for k in model_state.keys()):
        model_state = {k.replace('_orig_mod.', ''): v for k, v in model_state.items()}
    model.load_state_dict(model_state)

    if optimizer is not None:
        optimizer.load_state_dict(state_dict['optimizer_state_dict'])
    return state_dict['iteration']

def generate_stream(
    model, 
    tokenizer, 
    prompt,
    max_tokens=256,
    temperature=1.0,
    top_p=0.9,
    top_k=0,
    device='cuda',
    eos_token_id=None
):
    """
    流式生成文本,逐 token yield
    
    Yields:
        str: 每次生成的单个 token 解码后的文本
    """
    model.eval()
    
    input_ids = tokenizer.encode(prompt)
    input_ids = torch.tensor([input_ids], dtype=torch.long, device=device)
    # input_ids shape: (1, seq_len)
    
    if eos_token_id is None:
        eos_token_id = tokenizer.vocab_to_id.get(b"<|endoftext|>", None)
    
    with torch.no_grad():
        for _ in range(max_tokens):
            context_length = getattr(model, 'context_length', 256)
            # 截断到模型的 context_length
            context = input_ids[:, -context_length:]
            
            logits = model(context)
            next_token = choose_next_token(logits, temperature, top_p, top_k)
            
            if eos_token_id is not None and next_token == eos_token_id:
                break
            
            # 实时解码并 yield 单个 token
            token_text = tokenizer.decode([next_token])
            yield token_text
            
            next_token_tensor = torch.tensor([[next_token]], dtype=torch.long, device=device)
            input_ids = torch.cat([input_ids, next_token_tensor], dim=1)

def choose_next_token(logits, temperature, top_p, top_k):
    """
    temperature: 温度参数，控制采样的随机性
    top_p: 累积概率阈值
    top_k: top-k 采样的 k 值
    先top-k，再top-p
    使用torch.multinomial进行采样
    """
    next_token_logits = logits.squeeze()[-1]
    # next_token_logits.shape = (10000,)
    # 这是第 256 个位置对应的 10000 个词的 logits

    # 如果温度为0，直接返回最大概率的token（贪心解码）索引
    if temperature == 0:
        return torch.argmax(next_token_logits).item()
    
    # 应用温度缩放，如果 temperature > 1，分布更平坦，更创造；如果 < 1，分布更陡峭，更保守
    next_token_logits = next_token_logits / temperature
    
    # Top-k 采样：只保留概率最高的k个token
    if top_k > 0:
        # 获取top-k的值和索引
        top_k_values, top_k_indices = torch.topk(next_token_logits, min(top_k, next_token_logits.size(-1)))
        # 创建一个mask，只保留top-k的logits
        logits_mask = torch.full_like(next_token_logits, float('-inf'))
        logits_mask[top_k_indices] = top_k_values
        next_token_logits = logits_mask
    
    # 计算概率分布
    probabilities = torch.softmax(next_token_logits, dim=0)
    
    # Top-p 采样（Nucleus Sampling）
    if top_p < 1.0:
        # 按概率从大到小排序
        sorted_probs, sorted_indices = torch.sort(probabilities, descending=True)
        # 计算累积概率
        cumulative_probs = torch.cumsum(sorted_probs, dim=0)
        # 找到累积概率超过top_p的位置
        cutoff_index = torch.where(cumulative_probs > top_p)[0]
        if len(cutoff_index) > 0:
            # 保留累积概率在top_p范围内的token
            cutoff_index = cutoff_index[0].item() # 第一个超过top_p的索引
            # 创建新的概率分布，只保留累积概率在top_p范围内的token
            nucleus_probs = torch.zeros_like(probabilities)
            # 将累积概率在top_p范围内的token的概率赋值给新的概率分布
            nucleus_probs[sorted_indices[:cutoff_index + 1]] = sorted_probs[:cutoff_index + 1]
            # 重新归一化
            nucleus_probs = nucleus_probs / nucleus_probs.sum()
            probabilities = nucleus_probs
    
    # 从概率分布中采样
    next_token = torch.multinomial(probabilities, num_samples=1).item()
    
    return next_token

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        # 强制转换为 long (int64) 类型，防止传入 float 导致报错
        token_ids = token_ids.to(torch.long)

    # 【关键修改】强制转换为 long (int64)
    targets=targets.view(-1).to(torch.long)

        # Gather the cached tables for the required positions
        # [Batch, Seq] -> [Batch, Seq, Dim] -> unsqueeze -> [Batch, 1, Seq, Dim]
        cos_pos = self.cos_cached[token_positions].unsqueeze(1)
        sin_pos = self.sin_cached[token_positions].unsqueeze(1)

    # 提取目标类别的 logit，输出形状与 index 相同
    target_logits=logits.gather(dim=-1,index=targets.unsqueeze(1))

    # 【关键修复】将 (N, 1) 变为 (N)，防止广播成 (N, N)
    target_logits = target_logits.squeeze(-1)
    loss=logexpsum-target_logits

torch.set_float32_matmul_precision('high') # TF-32
model=torch.compile(model)