| import os | |
| import torch | |
| MODE = os.environ.get("TRANSLIT_MODE", "sentence") # "sentence" | "ime" | |
| class Config: | |
| # ---------------- common ---------------- | |
| d_model = 384 | |
| nhead = 6 | |
| num_layers = 4 | |
| dim_ff = 1536 | |
| dropout = 0.1 | |
| batch_size = 1024 | |
| lr = 7e-4 | |
| warmup_steps = 2000 | |
| weight_decay = 0.01 | |
| label_smooth = 0.1 | |
| grad_clip = 1.0 | |
| num_workers = 24 | |
| amp_dtype = torch.bfloat16 | |
| max_src_len = 256 | |
| max_tgt_len = 320 | |
| dry_run_samples = 20_000 | |
| dry_run_steps = 100 | |
| # ---------------- mode specific ---------------- | |
| if MODE == "ime": | |
| train_path = "data/word_train.jsonl" | |
| val_path = "data/word_val.jsonl" | |
| test_path = "data/word_test.jsonl" | |
| ckpt_dir = "checkpoints_ime" | |
| max_len = 64 | |
| batch_size = 1024 # big dataset again -> big batch | |
| epochs = 10 # millions of words -> few epochs suffice | |
| lr = 7e-4 | |
| warmup_steps = 2000 | |
| dropout = 0.1 # plenty of data -> normal dropout | |
| else: | |
| # sentence-level Malayalam -> Manglish (original) | |
| train_path = "data/train.jsonl" | |
| val_path = "data/val.jsonl" | |
| test_path = "data/test.jsonl" | |
| ckpt_dir = "checkpoints" | |
| max_len = 512 | |
| epochs = 8 | |
| CFG = Config() | |
| print(f"[config] mode={MODE} data={CFG.train_path} ckpt={CFG.ckpt_dir}") |