amalia 0.2.0__py3-none-any.whl

amalia-0.2.0.dist-info/METADATA

@@ -0,0 +1,31 @@
+Metadata-Version: 2.4
+Name: amalia
+Version: 0.2.0
+Summary: Add your description here
+Requires-Python: >=3.14
+Requires-Dist: numpy>=2.5.0
+Requires-Dist: torch>=2.12.1
+Description-Content-Type: text/markdown
+
+# AMALIA
+
+AMALIA is a decoder-only transformer architecture inherited from EuroLLM-9B, implemented in plain PyTorch (no external attention kernels).
+
+## Usage in Google Colab
+
+```python
+!pip install amalia
+
+import torch
+from amalia_core import AmaliaConfig, AmaliaForCausalLM
+
+# Initialize the model with random weights
+config = AmaliaConfig()
+model = AmaliaForCausalLM(config)
+
+# Run a forward pass on random token ids
+input_ids = torch.randint(0, config.vocab_size, (1, 16))
+logits = model(input_ids)
+
+print(logits.shape)  # torch.Size([1, 16, 128000])
+```

amalia-0.2.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+amalia_core/__init__.py,sha256=un23WcEPxoRbJguj7UV_ID0C29H_aPw42-V7XKXXpug,126
+amalia_core/architecture.py,sha256=G_bELz6Tz3_wjDgU0qfR2WZReJE46_RYpC4VKCk74LQ,5434
+amalia_core/config.py,sha256=5m8KXw7WMoLt0LHJS0dxx3cYyiCZG-CiPnvEKhqJNiA,513
+amalia-0.2.0.dist-info/METADATA,sha256=1aGOKUuyo7jCVSegurjfTiihPO-vIuBZIgPW57cR3OQ,766
+amalia-0.2.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+amalia-0.2.0.dist-info/RECORD,,

amalia-0.2.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

amalia_core/__init__.py

@@ -0,0 +1,4 @@
+from .config import AmaliaConfig
+from .architecture import AmaliaForCausalLM
+
+__all__ = ["AmaliaConfig", "AmaliaForCausalLM"]

amalia_core/architecture.py

@@ -0,0 +1,134 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import AmaliaConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return self.weight * x
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, head_dim: int, theta: float):
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, seq_len: int, device, dtype):
+        positions = torch.arange(seq_len, device=device, dtype=torch.float32)
+        freqs = torch.outer(positions, self.inv_freq.to(device))
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos().to(dtype), emb.sin().to(dtype)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    cos = cos[None, None, :, :]
+    sin = sin[None, None, :, :]
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    batch, n_kv_heads, seq_len, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    x = x[:, :, None, :, :].expand(batch, n_kv_heads, n_rep, seq_len, head_dim)
+    return x.reshape(batch, n_kv_heads * n_rep, seq_len, head_dim)
+
+
+class GroupedQueryAttention(nn.Module):
+    def __init__(self, config: AmaliaConfig):
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.head_dim = config.head_dim
+
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        batch, seq_len, _ = x.shape
+
+        q = self.q_proj(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        k = repeat_kv(k, self.num_kv_groups)
+        v = repeat_kv(v, self.num_kv_groups)
+
+        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+
+        attn_out = attn_out.transpose(1, 2).reshape(batch, seq_len, self.num_heads * self.head_dim)
+        return self.o_proj(attn_out)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: AmaliaConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: AmaliaConfig):
+        super().__init__()
+        self.input_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.attn = GroupedQueryAttention(config)
+        self.post_attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mlp = SwiGLU(config)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.input_norm(x), cos, sin)
+        x = x + self.mlp(self.post_attn_norm(x))
+        return x
+
+
+class AmaliaModel(nn.Module):
+    def __init__(self, config: AmaliaConfig):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.rotary_emb = RotaryEmbedding(config.head_dim, config.rope_theta)
+        self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        x = self.embed_tokens(input_ids)
+        cos, sin = self.rotary_emb(input_ids.shape[1], x.device, x.dtype)
+        for layer in self.layers:
+            x = layer(x, cos, sin)
+        return self.norm(x)
+
+
+class AmaliaForCausalLM(nn.Module):
+    def __init__(self, config: AmaliaConfig):
+        super().__init__()
+        self.model = AmaliaModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.model(input_ids)
+        return self.lm_head(hidden_states)

amalia_core/config.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class AmaliaConfig:
+    vocab_size: int = 128_000
+    hidden_size: int = 4096
+    intermediate_size: int = 12_288
+    num_hidden_layers: int = 42
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    max_position_embeddings: int = 32_768
+    rope_theta: float = 1_000_000.0
+    rms_norm_eps: float = 1e-5
+    tie_word_embeddings: bool = False
+
+    @property
+    def head_dim(self) -> int:
+        return self.hidden_size // self.num_attention_heads