47 lines
1.2 KiB
Python
47 lines
1.2 KiB
Python
"""Minimal CausalLM skeleton for the course implementation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
@dataclass
|
|
class CourseMiniMindConfig:
|
|
vocab_size: int = 6400
|
|
hidden_size: int = 768
|
|
num_hidden_layers: int = 8
|
|
num_attention_heads: int = 8
|
|
num_key_value_heads: int = 4
|
|
head_dim: int = 96
|
|
intermediate_size: int = 2432
|
|
max_position_embeddings: int = 32768
|
|
rope_theta: float = 1e6
|
|
rms_norm_eps: float = 1e-6
|
|
dropout: float = 0.0
|
|
|
|
|
|
class CourseMiniMindForCausalLM(nn.Module):
|
|
"""TODO: assemble embedding, blocks, final norm, lm_head, and loss.
|
|
|
|
Align with:
|
|
- model/model_minimind.py::MiniMindModel
|
|
- model/model_minimind.py::MiniMindForCausalLM
|
|
"""
|
|
|
|
def __init__(self, config: CourseMiniMindConfig):
|
|
super().__init__()
|
|
self.config = config
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: torch.Tensor,
|
|
attention_mask: torch.Tensor | None = None,
|
|
past_key_values: list[tuple[torch.Tensor, torch.Tensor] | None] | None = None,
|
|
use_cache: bool = False,
|
|
labels: torch.Tensor | None = None,
|
|
):
|
|
raise NotImplementedError("Implement during model assembly and pretrain lessons.")
|