Hugging Face's logo

mkd-hossain
/
finetune-model

Model card Files
xet
 Community
finetune-model / python_scripts /mkd_core.py
Md Najmul Hossain
Upload python_scripts/mkd_core.py with huggingface_hub
aaf22a7 verified
Raw History Blame Contribute Delete
4.77 kB
# ---------- mkd_core.py ----------
# (์ฃผ์‹ํšŒ์‚ฌ์— ์ผ€์ด๋”” | MKD)
# MKDCore Transformer โ€” Optimized for 10GB GPU
# Author: MKD AI Systems
import warnings
warnings.filterwarnings("ignore", category=FutureWarning)
import torch
from torch import nn
from tokenizers import Tokenizer
# ============================================================
# CONFIGURATION
# ============================================================
SEQ_LEN = 1024 # Context window (tokens)
D_MODEL = 1024 # Hidden dimension
N_HEADS = 16 # Attention heads
N_LAYERS = 12 # Transformer depth
FF_DIM = D_MODEL * 4 # Feed-forward expansion
DROPOUT = 0.1 # Regularization rate
# ============================================================
# ๐Ÿ”ค TOKENIZER WRAPPER
# ============================================================
class MyTokenizer:
"""
Wrapper for your trained BPE tokenizer (prod_tokenizer.json)
Adds <|im_start|> and <|im_end|> for chat/sequence training.
"""
def __init__(self, json_path: str):
self.tok = Tokenizer.from_file(json_path)
def encode(self, text: str, add_special: bool = True) -> list[int]:
if add_special:
text = "<|im_start|>" + text + "<|im_end|>"
return self.tok.encode(text).ids
def decode(self, ids: list[int]) -> str:
return self.tok.decode(ids, skip_special_tokens=True)
@property
def vocab_size(self) -> int:
return self.tok.get_vocab_size()
# ============================================================
# MODEL COMPONENTS
# ============================================================
class FeedForward(nn.Module):
"""Feed-Forward Network (FFN) with GELU activation"""
def __init__(self, dim: int):
super().__init__()
self.net = nn.Sequential(
nn.Linear(dim, FF_DIM),
nn.GELU(),
nn.Linear(FF_DIM, dim),
nn.Dropout(DROPOUT)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.net(x)
class MultiHeadAttention(nn.Module):
"""Scaled Dot-Product Multi-Head Attention"""
def __init__(self, dim: int, heads: int):
super().__init__()
self.heads = heads
self.scale = (dim // heads) ** -0.5
self.qkv = nn.Linear(dim, dim * 3)
self.proj = nn.Linear(dim, dim)
self.dropout = nn.Dropout(DROPOUT)
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, T, C = x.shape
qkv = self.qkv(x).chunk(3, dim=-1)
q, k, v = [t.view(B, T, self.heads, C // self.heads).transpose(1, 2) for t in qkv]
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.dropout(attn)
out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)
return self.proj(out)
class Block(nn.Module):
"""Transformer block: Attention + FeedForward"""
def __init__(self, dim: int, heads: int):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = MultiHeadAttention(dim, heads)
self.norm2 = nn.LayerNorm(dim)
self.ff = FeedForward(dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x + self.attn(self.norm1(x))
x = x + self.ff(self.norm2(x))
return x
# ============================================================
# ๐Ÿง  MKDCore TRANSFORMER MODEL
# ============================================================
class MKDCore(nn.Module):
"""MKDCore Transformer Language Model"""
def __init__(self, vocab_size: int):
super().__init__()
self.token_emb = nn.Embedding(vocab_size, D_MODEL)
self.pos_emb = nn.Embedding(SEQ_LEN, D_MODEL)
self.drop = nn.Dropout(DROPOUT)
self.layers = nn.ModuleList([Block(D_MODEL, N_HEADS) for _ in range(N_LAYERS)])
self.norm = nn.LayerNorm(D_MODEL)
self.head = nn.Linear(D_MODEL, vocab_size, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, T = x.shape
pos = torch.arange(0, T, device=x.device).unsqueeze(0)
x = self.token_emb(x) + self.pos_emb(pos)
x = self.drop(x)
for layer in self.layers:
x = layer(x)
x = self.norm(x)
logits = self.head(x)
return logits
# ============================================================
# โœ… TEST (Run only if executed directly)
# ============================================================
if __name__ == "__main__":
print("๐Ÿ” Testing MKDCore model...")
vocab_size = 32000
model = MKDCore(vocab_size=vocab_size)
dummy = torch.randint(0, vocab_size, (2, 64)) # batch=2, seq=64
out = model(dummy)
print(f"โœ… Output shape: {out.shape}") # [2, 64, 32000]