PredictiveManish/Trimurti-LM

1"""
2Step 2: Model configuration
3"""
4
5from dataclasses import dataclass
6from transformers import GPT2Config
7
8@dataclass
9class ModelConfig:
10    # Model architecture
11    vocab_size: int = 8000  # Updated from tokenizer
12    n_positions: int = 256  # Context length
13    n_embd: int = 512       # Hidden size
14    n_layer: int = 8        # Number of layers
15    n_head: int = 8         # Attention heads
16    n_inner: int = 1024     # FFN dimension
17    
18    # Training - REALISTIC VALUES
19    batch_size: int = 8     # Per GPU batch size
20    gradient_accumulation: int = 4  # Effective batch = 32
21    learning_rate: float = 3e-4
22    warmup_steps: int = 1000
23    total_steps: int = 20000  # ~8-9 epochs, NOT 50000
24    weight_decay: float = 0.1
25    max_grad_norm: float = 1.0
26    
27    # Data
28    train_file: str = "./final_corpus/multilingual_corpus_train.txt"
29    val_file: str = "./final_corpus/multilingual_corpus_val.txt"
30    tokenizer_path: str = "./final_corpus/multilingual_spm.model"
31    
32    # Checkpoints
33    output_dir: str = "./checkpoints"
34    save_steps: int = 1000
35    eval_steps: int = 500
36    logging_steps: int = 100
37    
38    # Mixed precision
39    fp16: bool = True
40    
41    def __post_init__(self):
42        print(f"\nModel Configuration (REALISTIC):")
43        print(f"  Parameters: ~{self.total_params:.1f}M")
44        print(f"  Hidden size: {self.n_embd}")
45        print(f"  Layers: {self.n_layer}")
46        print(f"  Context length: {self.n_positions}")
47        print(f"  Effective batch: {self.effective_batch_size}")
48        print(f"  Total steps: {self.total_steps} (~8-9 epochs)")
49        print(f"  Learning rate: {self.learning_rate}")
50    
51    @property
52    def effective_batch_size(self):
53        return self.batch_size * self.gradient_accumulation
54    
55    @property
56    def total_params(self):
57        # Rough estimate
58        embedding = self.vocab_size * self.n_embd
59        attention = 4 * self.n_embd * self.n_embd
60        ffn = 2 * self.n_embd * self.n_inner
61        ln = 2 * self.n_embd
62        per_layer = attention + ffn + ln
63        total = embedding + (self.n_layer * per_layer)
64        return total / 1e6  # Millions