Update Mo_jax.py

Mo_jax.py

@@ -8,7 +8,7 @@ import numpy as np
 import sentencepiece as spm
 from functools import partial
 from typing import Any, Callable, Optional, Tuple, Sequence
-
+import requests
 import jax
 import jax.numpy as jnp
 from jax import random
@@ -17,6 +17,13 @@ from flax.training import train_state, checkpoints
 import optax
 import tqdm
 
+def download_file(url, save_path):
+    r = requests.get(url, stream=True)
+    r.raise_for_status()
+    with open(save_path, "wb") as f:
+        for chunk in r.iter_content(8192*2):
+            f.write(chunk)
+    print(f"✅ {save_path} 저장됨")
 # ------------------
 # Config
 # ------------------
@@ -32,6 +39,18 @@ SEED = 42
 LEARNING_RATE = 1e-4
 EPOCHS = 1
 
+if not os.path.exists(CORPUS_PATH):
+    download_file(
+        "https://huggingface.co/datasets/Yuchan5386/Prototype/resolve/main/corpus_ko.txt?download=true",
+        CORPUS_PATH
+    )
+
+if not os.path.exists(VOCAB_MODEL):
+    download_file(
+        "https://huggingface.co/Yuchan5386/inlam-100m/resolve/main/ko_unigram.model?download=true",
+        VOCAB_MODEL
+    )
+
 # Derived
 NUM_DEVICES = jax.device_count()
 assert GLOBAL_BATCH % NUM_DEVICES == 0, "GLOBAL_BATCH must be divisible by device count"
@@ -379,3 +398,4 @@ def generate_text(state, prompt: str, max_gen=256, p=0.9, temperature=0.8, min_l
 # quick generate
 print("\n\n===== 생성 결과 =====")
 print(generate_text(state, "지난 2년 동안 출연연이 국가가 필요한 연구를", p=0.9))
+