Upload train_autogrow.py with huggingface_hub

train_autogrow.py

@@ -99,7 +99,8 @@ def run_test(model, tokenizer, step):
         norm_r = F.normalize(r_noise, dim=-1)
         logits = torch.matmul(norm_r, norm_weights.T)
         resp_ids = torch.argmax(logits, dim=-1)
-        result = tokenizer.decode(resp_ids[0], skip_special_tokens=True)
+        # Show special tokens to debug why it's silent
+        result = tokenizer.decode(resp_ids[0], skip_special_tokens=False)
         log(f"Prompt: '{prompt}' | [Cropmark]: '{result}'")
         with open(CHECKPOINT_LOG, "a") as f:
             f.write(f"Step {step} - Prompt: '{prompt}' | [Cropmark]: '{result}'\n")
@@ -130,13 +131,22 @@ if __name__ == "__main__":
         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:])
+            prompt_ids = input_ids[:, :MAX_PROMPT_LEN]
+            resp_ids = input_ids[:, MAX_PROMPT_LEN:]
+            
+            prompt_emb = model.token_embedding(prompt_ids)
+            resp_emb = model.token_embedding(resp_ids)
             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()
+            # Cosine Similarity Loss with Padding Mask
+            mask = (resp_ids != tokenizer.pad_token_id).float()
+            # Calculate cosine similarity for each token
+            cos_sim = F.cosine_similarity(pred_resp, resp_emb, dim=-1)
+            # Mask out padding tokens
+            loss = 1 - (cos_sim * mask).sum() / (mask.sum() + 1e-8)
+            
             loss.backward()
             optimizer.step()
             if step % 100 == 0: