Views
No views yet
Model bahasa Indonesia eksperimental berbasis arsitektur Mamba2 (State Space Model).Experimental Indonesian language model based on Mamba2 (State Space Model) architecture.
| Detail | Value |
|---|---|
| Architecture | Mamba2 + RMSNorm |
| Parameters | 38 Million |
| Vocab Size | 30,521 (IndoBERT) |
| Training Data | Indonesian Wikipedia (subset) |
| Training Steps | 3 Epochs 3.8% of real dataset |
| Max Sequence | 256 tokens |
| Precision | FP16 (AMP) |
1import torch
2from transformers import AutoTokenizer
3from mamba_ssm import Mamba2
4
5# 1. Load Tokenizer
6tokenizer = AutoTokenizer.from_pretrained("indobenchmark/indobert-base-p1")
7tokenizer.pad_token = tokenizer.eos_token
8
9# 2. Define Model Class (same as training) / Definisi Class Model (sama seperti training)
10class RMSNorm(torch.nn.Module):
11 def __init__(self, d_model, eps=1e-5):
12 super().__init__()
13 self.eps = eps
14 self.weight = torch.nn.Parameter(torch.ones(d_model))
15
16 def forward(self, x):
17 norm = x.pow(2).mean(-1, keepdim=True)
18 x_normed = x * torch.rsqrt(norm + self.eps)
19 return self.weight * x_normed
20
21class OtterMambaLM(torch.nn.Module):
22 def __init__(self, vocab_size, d_model=768, n_layer=4, max_seq_len=256):
23 super().__init__()
24 self.embedding = torch.nn.Embedding(vocab_size, d_model)
25 self.layers = torch.nn.ModuleList([
26 Mamba2(d_model=d_model, d_state=64, d_conv=4, expand=2)
27 for _ in range(n_layer)
28 ])
29 self.norm_f = RMSNorm(d_model)
30 self.lm_head = torch.nn.Linear(d_model, vocab_size, bias=False)
31 self.lm_head.weight = self.embedding.weight # Weight tying
32 self.pos_emb = torch.nn.Parameter(torch.zeros(1, max_seq_len, d_model))
33
34 def forward(self, input_ids):
35 x = self.embedding(input_ids) + self.pos_emb[:, :input_ids.shape[1], :]
36 for layer in self.layers:
37 out = layer(x)
38 if isinstance(out, tuple): out = out[0] # Handle tuple output
39 x = x + out # Residual connection
40 x = self.norm_f(x)
41 return self.lm_head(x)
42
43# 3. Load Weights / Muat Bobot
44model = OtterMambaLM(vocab_size=30521, d_model=768, n_layer=4, max_seq_len=256)
45model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
46model.eval()
47
48# 4. Generate Text / Hasilkan Teks
49input_ids = tokenizer.encode("Indonesia adalah", return_tensors="pt")
50with torch.no_grad():
51 logits = model(input_ids)
52 next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
53 output = tokenizer.decode(next_token[0])
54 print(output)
55
56## ⚠️ Limitations & Known Issues / Keterbatasan