| import os |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from tokenizers import Tokenizer, pre_tokenizers, decoders |
|
|
| MODEL_PATH = "model.pt" |
| TOKENIZER_PATH = "tokenizer.json" |
|
|
| D_MODEL = 512 |
| NUM_HEADS = 8 |
| NUM_KV_HEADS = 8 |
| NUM_LAYERS = 18 |
| D_FF = 2048 |
| MAX_SEQ_LENGTH = 768 |
| VOCAB_SIZE_LIMIT = 32768 |
| USE_LAYER_SCALE = False |
|
|
| class BPETokenizer: |
| def __init__(self, vocab_size=15000): |
| self.tokenizer = None |
| self.vocab = {} |
| self.token_to_id = {} |
| self.id_to_token = {} |
|
|
| def load(self, filepath): |
| self.tokenizer = Tokenizer.from_file(filepath) |
| self.tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) |
| self.tokenizer.decoder = decoders.ByteLevel() |
| |
| self.vocab = self.tokenizer.get_vocab() |
| self.token_to_id = self.vocab |
| self.id_to_token = {v: k for k, v in self.vocab.items()} |
| print(f"Tokenizer loaded. Vocabulary size: {self.tokenizer.get_vocab_size()}") |
|
|
| def encode(self, text): |
| return self.tokenizer.encode(text).ids |
|
|
| def decode(self, ids): |
| return self.tokenizer.decode(ids, skip_special_tokens=False) |
|
|
| class FastRMSNorm(nn.Module): |
| def __init__(self, d_model, eps=1e-5): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(d_model)) |
|
|
| def forward(self, x): |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight |
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__(self, dim, max_seq_len=4096, theta=10000.0): |
| super().__init__() |
| self.dim = dim |
| self.theta = theta |
| self.max_seq_len = max_seq_len |
| |
| inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer("inv_freq", inv_freq) |
| self._update_cache(max_seq_len) |
|
|
| def _update_cache(self, seq_len): |
| self.max_seq_len = seq_len |
| t = torch.arange(self.max_seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype) |
| freqs = torch.outer(t, self.inv_freq) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| |
| self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False) |
| self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False) |
|
|
| def forward(self, x, seq_len): |
| if seq_len > self.max_seq_len: |
| self._update_cache(seq_len) |
| return ( |
| self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype), |
| self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype) |
| ) |
|
|
| def apply_rotary_pos_emb(x, cos, sin): |
| def rotate_half(x): |
| x1, x2 = x.chunk(2, dim=-1) |
| return torch.cat((-x2, x1), dim=-1) |
| return (x * cos) + (rotate_half(x) * sin) |
|
|
| def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: |
| batch, n_kv_heads, seqlen, head_dim = x.shape |
| if n_rep == 1: |
| return x |
| return ( |
| x[:, :, None, :, :] |
| .expand(batch, n_kv_heads, n_rep, seqlen, head_dim) |
| .reshape(batch, n_kv_heads * n_rep, seqlen, head_dim) |
| ) |
|
|
| def apply_swiglu(gate_up): |
| a, b = gate_up.chunk(2, dim=-1) |
| return F.silu(a) * b |
|
|
| class TokenEmbedding(nn.Embedding): |
| def __init__(self, vocab_size, d_model, pad_id=None): |
| padding_idx = pad_id if pad_id is not None and pad_id >= 0 else None |
| super().__init__(num_embeddings=vocab_size, embedding_dim=d_model, padding_idx=padding_idx) |
|
|
| class GroupedQueryAttention(nn.Module): |
| def __init__(self, d_model, num_heads, num_kv_heads): |
| super().__init__() |
| self.d_model = d_model |
| self.num_heads = num_heads |
| self.num_kv_heads = num_kv_heads |
| self.head_dim = d_model // num_heads |
| self.n_rep = self.num_heads // self.num_kv_heads |
|
|
| self.W_q = nn.Linear(d_model, num_heads * self.head_dim, bias=False) |
| self.W_k = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False) |
| self.W_v = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False) |
| self.W_o = nn.Linear(num_heads * self.head_dim, d_model, bias=False) |
| |
| self.rope = RotaryEmbedding(self.head_dim) |
| |
| self.q_norm = FastRMSNorm(self.head_dim) |
| self.k_norm = FastRMSNorm(self.head_dim) |
|
|
| def forward(self, q_in, k_in, v_in): |
| batch_size, seq_len, _ = q_in.shape |
| |
| q = self.W_q(q_in) |
| k = self.W_k(k_in) |
| v = self.W_v(v_in) |
| |
| q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) |
| k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) |
| v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) |
| |
| q = self.q_norm(q) |
| k = self.k_norm(k) |
|
|
| q = q.transpose(1, 2) |
| k = k.transpose(1, 2) |
| v = v.transpose(1, 2) |
| |
| cos_q, sin_q = self.rope(q, q.shape[2]) |
| cos_k, sin_k = self.rope(k, k.shape[2]) |
| |
| q = apply_rotary_pos_emb(q, cos_q, sin_q) |
| k = apply_rotary_pos_emb(k, cos_k, sin_k) |
|
|
| k = repeat_kv(k, self.n_rep) |
| v = repeat_kv(v, self.n_rep) |
|
|
| context = F.scaled_dot_product_attention( |
| q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True |
| ) |
| context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) |
| return self.W_o(context) |
|
|
| class PositionWiseFeedForward(nn.Module): |
| def __init__(self, d_model, d_ff, bias=False): |
| super().__init__() |
| self.gate_up_proj = nn.Linear(d_model, d_ff * 2, bias=bias) |
| self.down_proj = nn.Linear(d_ff, d_model, bias=bias) |
|
|
| def forward(self, x): |
| gate_up = self.gate_up_proj(x) |
| activated = apply_swiglu(gate_up) |
| return self.down_proj(activated) |
|
|
| class DecoderLayer(nn.Module): |
| def __init__(self, d_model, num_heads, num_kv_heads, d_ff, use_layer_scale=False): |
| super().__init__() |
| self.norm1 = FastRMSNorm(d_model) |
| self.norm2 = FastRMSNorm(d_model) |
|
|
| self.attn = GroupedQueryAttention(d_model, num_heads, num_kv_heads) |
| self.ffn = PositionWiseFeedForward(d_model, d_ff) |
| |
| self.use_layer_scale = use_layer_scale |
| if self.use_layer_scale: |
| init_value = 1e-4 |
| self.ls_1 = nn.Parameter(torch.ones(d_model) * init_value) |
| self.ls_2 = nn.Parameter(torch.ones(d_model) * init_value) |
|
|
| def forward(self, x): |
| x_norm = self.norm1(x) |
| attn_out = self.attn(x_norm, x_norm, x_norm) |
| x = x + (attn_out * self.ls_1 if self.use_layer_scale else attn_out) |
| |
| x_norm = self.norm2(x) |
| ffn_out = self.ffn(x_norm) |
| x = x + (ffn_out * self.ls_2 if self.use_layer_scale else ffn_out) |
| |
| return x |
|
|
| class Transformer(nn.Module): |
| def __init__(self, vocab_size, d_model, num_heads, num_kv_heads, num_layers, d_ff, pad_id, use_layer_scale=False): |
| super().__init__() |
| self.pad_id = pad_id |
| self.embedding = TokenEmbedding(vocab_size, d_model, pad_id) |
| self.layers = nn.ModuleList([ |
| DecoderLayer(d_model, num_heads, num_kv_heads, d_ff, use_layer_scale) for _ in range(num_layers) |
| ]) |
| self.norm_f = FastRMSNorm(d_model) |
| self.fc_out = nn.Linear(d_model, vocab_size, bias=False) |
|
|
| def forward(self, x): |
| x = self.embedding(x) |
| for layer in self.layers: |
| x = layer(x) |
| x = self.norm_f(x) |
| return self.fc_out(x) |
|
|
| @torch.no_grad() |
| def generate_response(model, tokenizer, user_input, max_seq_length, device, temperature=0.7, top_k=50, repetition_penalty=1.2): |
| model.eval() |
|
|
| SOS_ID = tokenizer.token_to_id.get("<sos>", None) |
| EOS_ID = tokenizer.token_to_id.get("<eos>", None) |
|
|
| if not user_input.strip(): |
| return "Please say something." |
|
|
| input_ids = tokenizer.encode(user_input) |
| if SOS_ID is not None: |
| input_ids = [SOS_ID] + input_ids |
|
|
| generated_tokens_set = set() |
| generated_new_tokens = [] |
|
|
| for _ in range(max_seq_length - len(input_ids)): |
| x_tensor = torch.LongTensor([input_ids]).to(device) |
| output = model(x_tensor) |
| |
| last_logits = output[0, -1, :] / temperature |
|
|
| for token_id in generated_tokens_set: |
| if last_logits[token_id] > 0: |
| last_logits[token_id] /= repetition_penalty |
| else: |
| last_logits[token_id] *= repetition_penalty |
| |
| if top_k > 0: |
| v, _ = torch.topk(last_logits, top_k) |
| last_logits[last_logits < v[-1]] = -float('Inf') |
|
|
| probs = torch.softmax(last_logits, dim=-1) |
| next_word_id = torch.multinomial(probs, num_samples=1).item() |
| |
| if next_word_id == EOS_ID: |
| break |
| |
| input_ids.append(next_word_id) |
| generated_new_tokens.append(next_word_id) |
| generated_tokens_set.add(next_word_id) |
| |
| raw_response = tokenizer.decode(generated_new_tokens) |
|
|
| for st in ["<sos>", "<eos>", "<pad>"]: |
| raw_response = raw_response.replace(st, "") |
| |
| return raw_response.strip().capitalize() |
|
|
| if __name__ == '__main__': |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| print(f"Using device: {device}") |
|
|
| if not os.path.exists(TOKENIZER_PATH): |
| raise FileNotFoundError(f"Tokenizer not found at {TOKENIZER_PATH}") |
|
|
| tokenizer = BPETokenizer(vocab_size=VOCAB_SIZE_LIMIT) |
| tokenizer.load(TOKENIZER_PATH) |
|
|
| vocab_size = len(tokenizer.vocab) |
| PAD_ID = tokenizer.token_to_id.get("<pad>", 0) |
|
|
| print("Initializing model architecture...") |
| model = Transformer( |
| vocab_size=vocab_size, |
| d_model=D_MODEL, |
| num_heads=NUM_HEADS, |
| num_kv_heads=NUM_KV_HEADS, |
| num_layers=NUM_LAYERS, |
| d_ff=D_FF, |
| pad_id=PAD_ID, |
| use_layer_scale=USE_LAYER_SCALE |
| ) |
|
|
| if not os.path.exists(MODEL_PATH): |
| raise FileNotFoundError(f"Model file not found: {MODEL_PATH}") |
| |
| print(f"Loading weights from {MODEL_PATH}...") |
| model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu")) |
| print("Weights loaded successfully!") |
|
|
| dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 |
| model = model.to(dtype=dtype).to(device) |
|
|
| print("Model ready! Type 'exit' to quit.\n") |
|
|
| while True: |
| try: |
| user_msg = input("You: ") |
| if user_msg.lower() in ['exit', 'quit']: |
| break |
| |
| response = generate_response( |
| model=model, |
| tokenizer=tokenizer, |
| user_input=user_msg, |
| max_seq_length=MAX_SEQ_LENGTH, |
| device=device |
| ) |
| print(f"Bot: {response}\n") |
| |
| except KeyboardInterrupt: |
| print("\nExiting...") |
| break |
