| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| IM_START = 151644 |
| IM_END = 151645 |
| EOT = 151643 |
| THINK_END = 151668 |
|
|
| EXOCORE_IDENTITY = ( |
| "You are ExocoreV1, a highly capable AI built from scratch by Exocore. " |
| "You are knowledgeable, analytical, direct, and precise. " |
| "You reason carefully before answering. " |
| "You identify yourself only as ExocoreV1." |
| ) |
|
|
|
|
| def build_chat_prompt(messages, think=False, prior_thinking=None, search_context=None): |
| identity = EXOCORE_IDENTITY |
| if search_context: |
| identity += f"\n\nResearch context:\n{search_context}" |
|
|
| prompt = f"<|im_start|>system\n{identity}<|im_end|>\n" |
| for msg in messages: |
| role = msg.get("role", "user") |
| content = msg.get("content", "") |
| prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n" |
|
|
| if prior_thinking: |
| prompt += f"<|im_start|>assistant\n<think>\n{prior_thinking}\n\n</think>\n" |
| elif think: |
| prompt += "<|im_start|>assistant\n<think>\n" |
| else: |
| prompt += "<|im_start|>assistant\n<think>\n\n</think>\n" |
| return prompt |
|
|
|
|
| class ExocoreRMSNorm(nn.Module): |
| def __init__(self, dim, eps=1e-6): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x): |
| xf = x.float() |
| norm = (xf.pow(2).mean(-1, keepdim=True) + self.eps).rsqrt() |
| return (xf * norm).to(x.dtype) * self.weight |
|
|
|
|
| class ExocoreRotaryEmbedding(nn.Module): |
| def __init__(self, dim, max_seq=40960, base=1000000): |
| super().__init__() |
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer("inv_freq", inv_freq, persistent=False) |
| t = torch.arange(max_seq).float() |
| freqs = torch.outer(t, inv_freq) |
| cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1) |
| sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1) |
| self.register_buffer("rope_cos", cos, persistent=False) |
| self.register_buffer("rope_sin", sin, persistent=False) |
|
|
| def apply_rope(self, q, k): |
| T = q.shape[2] |
| cos = self.rope_cos[:T].unsqueeze(0).unsqueeze(0) |
| sin = self.rope_sin[:T].unsqueeze(0).unsqueeze(0) |
|
|
| def rotate_half(x): |
| half = x.shape[-1] // 2 |
| return torch.cat([-x[..., half:], x[..., :half]], dim=-1) |
|
|
| q_rot = q * cos + rotate_half(q) * sin |
| k_rot = k * cos + rotate_half(k) * sin |
| return q_rot.to(q.dtype), k_rot.to(k.dtype) |
|
|
|
|
| class ExocoreAttention(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.n_heads = cfg["num_attention_heads"] |
| self.n_kv_heads = cfg["num_key_value_heads"] |
| self.head_dim = cfg["head_dim"] |
| hidden = cfg["hidden_size"] |
| eps = cfg["rms_norm_eps"] |
|
|
| self.q_proj = nn.Linear(hidden, self.n_heads * self.head_dim, bias=False) |
| self.k_proj = nn.Linear(hidden, self.n_kv_heads * self.head_dim, bias=False) |
| self.v_proj = nn.Linear(hidden, self.n_kv_heads * self.head_dim, bias=False) |
| self.o_proj = nn.Linear(self.n_heads * self.head_dim, hidden, bias=False) |
| self.q_norm = ExocoreRMSNorm(self.head_dim, eps=eps) |
| self.k_norm = ExocoreRMSNorm(self.head_dim, eps=eps) |
|
|
| def forward(self, x, rope, mask=None): |
| B, T, _ = x.shape |
| q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) |
| k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) |
| v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) |
|
|
| q = self.q_norm(q) |
| k = self.k_norm(k) |
| q, k = rope.apply_rope(q, k) |
|
|
| groups = self.n_heads // self.n_kv_heads |
| k = k.repeat_interleave(groups, dim=1) |
| v = v.repeat_interleave(groups, dim=1) |
|
|
| scale = self.head_dim ** -0.5 |
| attn = torch.matmul(q, k.transpose(-2, -1)) * scale |
| if mask is not None: |
| attn = attn + mask |
| attn = F.softmax(attn.float(), dim=-1).to(q.dtype) |
|
|
| out = torch.matmul(attn, v) |
| out = out.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim) |
| return self.o_proj(out) |
|
|
|
|
| class ExocoreMLP(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.gate_proj = nn.Linear(cfg["hidden_size"], cfg["intermediate_size"], bias=False) |
| self.up_proj = nn.Linear(cfg["hidden_size"], cfg["intermediate_size"], bias=False) |
| self.down_proj = nn.Linear(cfg["intermediate_size"], cfg["hidden_size"], bias=False) |
|
|
| def forward(self, x): |
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| class ExocoreBlock(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.input_layernorm = ExocoreRMSNorm(cfg["hidden_size"], eps=cfg["rms_norm_eps"]) |
| self.post_attention_layernorm = ExocoreRMSNorm(cfg["hidden_size"], eps=cfg["rms_norm_eps"]) |
| self.attn = ExocoreAttention(cfg) |
| self.mlp = ExocoreMLP(cfg) |
|
|
| def forward(self, x, rope, mask=None): |
| x = x + self.attn(self.input_layernorm(x), rope, mask) |
| x = x + self.mlp(self.post_attention_layernorm(x)) |
| return x |
|
|
|
|
| class ExocoreLM(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.cfg = cfg |
| self.embed_tokens = nn.Embedding(cfg["vocab_size"], cfg["hidden_size"]) |
| self.layers = nn.ModuleList([ExocoreBlock(cfg) for _ in range(cfg["num_hidden_layers"])]) |
| self.norm = ExocoreRMSNorm(cfg["hidden_size"], eps=cfg["rms_norm_eps"]) |
| self.lm_head = nn.Linear(cfg["hidden_size"], cfg["vocab_size"], bias=False) |
| self.rope = ExocoreRotaryEmbedding( |
| cfg["head_dim"], |
| max_seq=cfg.get("max_position_embeddings", 40960), |
| base=cfg.get("rope_theta", 1000000), |
| ) |
|
|
| def forward(self, input_ids): |
| B, T = input_ids.shape |
| x = self.embed_tokens(input_ids) |
| mask = torch.triu(torch.full((T, T), float("-inf"), device=x.device), diagonal=1) |
| mask = mask.unsqueeze(0).unsqueeze(0) |
| for block in self.layers: |
| x = block(x, self.rope, mask) |
| x = self.norm(x) |
| return self.lm_head(x), None |
|
|
