Upload train_autogrow.py with huggingface_hub

train_autogrow.py

@@ -12,10 +12,11 @@ from huggingface_hub import HfApi
 # --- FAILPROOF CONFIG ---
 MODEL_PATH = "./DiffReaper-Talk"
 REPO_ID = "darwinkernelpanic/DiffReaper-5"
-HF_TOKEN = "${HF_TOKEN}"
+HF_TOKEN = os.getenv("HF_TOKEN")
 OUTPUT_DIR = "./training_output"
 LOG_FILE = "training.log"
-BATCH_SIZE = 16 # Lower for 3090 VRAM
+CHECKPOINT_LOG = "checkpoint_log.txt"
+BATCH_SIZE = 32
 LEARNING_RATE = 1e-4
 SAVE_EVERY = 2500
 TEST_EVERY = 500
@@ -79,23 +80,7 @@ class DiffReaperModel(nn.Module):
         for block in self.blocks: x = block(x, t_emb)
         return self.ln_f(x)
 
-log("Initializing Autogrow Model...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
-
-model = DiffReaperModel(tokenizer.vocab_size, N_EMBD, N_HEAD, N_LAYER).to("cuda")
-noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="squaredcos_cap_v2")
-optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
-
-log("Loading Dataset...")
-dataset = load_dataset("OpenAssistant/oasst1", split="train")
-def tokenize_function(examples):
-    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=TOTAL_LEN)
-tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)
-tokenized_dataset.set_format("torch")
-dataloader = torch.utils.data.DataLoader(tokenized_dataset, batch_size=BATCH_SIZE, shuffle=True)
-
-def run_test(step):
+def run_test(model, tokenizer, step):
     log(f"Running Cropmark Diagnostic [Step {step}]...")
     model.eval()
     with torch.no_grad():
@@ -113,37 +98,56 @@ def run_test(step):
         norm_r = F.normalize(r_noise, dim=-1)
         logits = torch.matmul(norm_r, norm_weights.T)
         resp_ids = torch.argmax(logits, dim=-1)
-        log(f"Prompt: '{prompt}' | [Cropmark]: '{tokenizer.decode(resp_ids[0], skip_special_tokens=True)}'")
+        result = tokenizer.decode(resp_ids[0], skip_special_tokens=True)
+        log(f"Prompt: '{prompt}' | [Cropmark]: '{result}'")
+        with open(CHECKPOINT_LOG, "a") as f:
+            f.write(f"Step {step} - Prompt: '{prompt}' | [Cropmark]: '{result}'\n")
     model.train()
 
-log("Autonomous growth starting...")
-api = HfApi()
-start_time = time.time()
-step = 0
-while True: # Unlimited steps, controlled by your credit
-    for batch in dataloader:
-        optimizer.zero_grad()
-        input_ids = batch["input_ids"].to("cuda")
-        prompt_emb = model.token_embedding(input_ids[:, :MAX_PROMPT_LEN])
-        resp_emb = model.token_embedding(input_ids[:, MAX_PROMPT_LEN:])
-        
-        noise = torch.randn_like(resp_emb)
-        t = torch.randint(0, 1000, (input_ids.shape[0],), device="cuda").long()
-        noisy_resp = noise_scheduler.add_noise(resp_emb, noise, t)
-        
-        pred_resp = model(torch.cat([prompt_emb, noisy_resp], dim=1), t)[:, MAX_PROMPT_LEN:, :]
-        loss = 1 - F.cosine_similarity(pred_resp, resp_emb, dim=-1).mean()
-        loss.backward()
-        optimizer.step()
-        
-        if step % 100 == 0:
-            elapsed = time.time() - start_time
-            log(f"Step {step} - Loss: {loss.item():.6f} - Speed: {(step+1)/elapsed:.2f} s/s")
-        if step > 0 and step % TEST_EVERY == 0: run_test(step)
-        if step > 0 and step % SAVE_EVERY == 0:
-            ckpt_path = os.path.join(OUTPUT_DIR, f"cropmark_latest.pt")
-            torch.save(model.state_dict(), ckpt_path)
-            log("Syncing to HF...")
-            try: api.upload_file(path_or_fileobj=ckpt_path, path_in_repo="cropmark_latest.pt", repo_id=REPO_ID, token=HF_TOKEN)
-            except Exception as e: log(f"HF Sync Error: {e}")
-        step += 1
+if __name__ == "__main__":
+    log("Initializing Autogrow Model...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
+
+    model = DiffReaperModel(tokenizer.vocab_size, N_EMBD, N_HEAD, N_LAYER).to("cuda")
+    noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="squaredcos_cap_v2")
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
+
+    log("Loading Dataset...")
+    dataset = load_dataset("OpenAssistant/oasst1", split="train")
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=TOTAL_LEN)
+    tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)
+    tokenized_dataset.set_format("torch")
+    dataloader = torch.utils.data.DataLoader(tokenized_dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+    log("Autonomous growth starting...")
+    api = HfApi()
+    start_time = time.time()
+    step = 0
+    while True:
+        for batch in dataloader:
+            optimizer.zero_grad()
+            input_ids = batch["input_ids"].to("cuda")
+            prompt_emb = model.token_embedding(input_ids[:, :MAX_PROMPT_LEN])
+            resp_emb = model.token_embedding(input_ids[:, MAX_PROMPT_LEN:])
+            noise = torch.randn_like(resp_emb)
+            t = torch.randint(0, 1000, (input_ids.shape[0],), device="cuda").long()
+            noisy_resp = noise_scheduler.add_noise(resp_emb, noise, t)
+            pred_resp = model(torch.cat([prompt_emb, noisy_resp], dim=1), t)[:, MAX_PROMPT_LEN:, :]
+            loss = 1 - F.cosine_similarity(pred_resp, resp_emb, dim=-1).mean()
+            loss.backward()
+            optimizer.step()
+            if step % 100 == 0:
+                elapsed = time.time() - start_time
+                log(f"Step {step} - Loss: {loss.item():.6f} - Speed: {(step+1)/elapsed:.2f} s/s")
+            if step > 0 and step % TEST_EVERY == 0: run_test(model, tokenizer, step)
+            if step > 0 and step % SAVE_EVERY == 0:
+                ckpt_path = os.path.join(OUTPUT_DIR, f"cropmark_{step}.pt")
+                torch.save(model.state_dict(), ckpt_path)
+                log("Syncing to HF...")
+                try: 
+                    api.upload_file(path_or_fileobj=ckpt_path, path_in_repo=f"cropmark_{step}.pt", repo_id=REPO_ID, token=HF_TOKEN)
+                    api.upload_file(path_or_fileobj="train_autogrow.py", path_in_repo="train_autogrow.py", repo_id=REPO_ID, token=HF_TOKEN)
+                except Exception as e: log(f"HF Sync Error: {e}")
+            step += 1