Upload train.py with huggingface_hub

train.py

@@ -61,7 +61,7 @@ def estimate_loss(model, data, args, device):
     return sum(losses) / len(losses)
 
 
-def main():
+def main(argv=None):
     parser = argparse.ArgumentParser(description="Latih model Indigo dari scratch")
     parser.add_argument("--data", nargs="+", default=["data/sample.txt"], help="path file/folder teks untuk training")
     parser.add_argument("--out", default="out", help="folder output checkpoint")
@@ -83,7 +83,7 @@ def main():
     parser.add_argument("--tokenizer", default="char", choices=["char", "bpe"])
     parser.add_argument("--vocab-size", type=int, default=512, help="ukuran vocab untuk tokenizer bpe")
     parser.add_argument("--val-fraction", type=float, default=0.1, help="proporsi file untuk validasi")
-    args = parser.parse_args()
+    args = parser.parse_args(argv)
 
     torch.manual_seed(args.seed)
     if args.device == "auto":
@@ -111,6 +111,7 @@ def main():
     start_step = 0
     init_meta = None
     config = None
+    comp_ratio = 1.0
     if args.init_from:
         init_state, init_meta, init_opt = load_init(args.init_from)
         config = GPTConfig(**init_meta["config"])
@@ -125,9 +126,10 @@ def main():
             tokenizer = BPETokenizer.train(all_text, args.vocab_size)
             tinfo = tokenizer.state()
             n_chars = len(all_text.encode("utf-8"))
+            comp_ratio = n_chars / max(1, len(tokenizer.encode(all_text)))
             print(
                 f"tokenizer=bpe | vocab={tokenizer.vocab_size} | "
-                f"kompresi {n_chars:,} karakter -> rasio {n_chars / max(1, len(tokenizer.encode(all_text))):.2f}x"
+                f"kompresi {n_chars:,} karakter -> rasio {comp_ratio:.2f}x"
             )
         else:
             tokenizer = CharTokenizer.from_text(all_text)
@@ -232,6 +234,35 @@ def main():
     torch.save(optimizer.state_dict(), os.path.join(args.out, "indigo_optimizer.pt"))
     print(f"model tersimpan di {final_path} (+_meta.json, indigo_optimizer.pt)")
 
+    stats = {
+        "out": args.out,
+        "device": device,
+        "backend": "pytorch",
+        "tokenizer": tinfo.get("type", "char"),
+        "vocab_size": tokenizer.vocab_size,
+        "compression_ratio": round(comp_ratio, 4),
+        "tokens_train": len(train_data),
+        "tokens_val": len(val_data),
+        "files_train": max(0, len(files) - n_val),
+        "files_val": n_val,
+        "steps_trained": args.steps,
+        "total_steps": total_steps,
+        "best_val": best_val if best_val != float("inf") else None,
+        "last_val": last_val,
+        "nats_per_char_best": (
+            round(best_val / comp_ratio, 4) if best_val != float("inf") and comp_ratio else None
+        ),
+        "params_million": round(model.num_params() / 1e6, 4),
+        "config": config.__dict__,
+        "args": {k: v for k, v in vars(args).items() if k != "data"},
+        "elapsed_sec": round(time.time() - t0, 1),
+    }
+    print(
+        f"ringkasan: best_val={stats['best_val']} | "
+        f"nats/karakter={stats['nats_per_char_best']} | params={stats['params_million']}M"
+    )
+    return stats
+
 
 if __name__ == "__main__":
     main()