Update copy_weights.py

copy_weights.py

@@ -1,5 +1,5 @@
 # =============================================================================
-# COPYRIGHT © 2025 Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED.
+# COPYRIGHT © 2025-2026 Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED.
 # CMS Manhattan JiRack Technology — PATENT PENDING
 #
 # This code is proprietary. 
@@ -18,24 +18,18 @@
 # 
 # =============================================================================
 
+
 import torch
 from JiRackTernary_new import JiRackConfig, JiRackTernary1B
 from transformers import AutoTokenizer
+from safetensors.torch import load_file, save_file
 import os
 
 print("🚀 Copying embeddings and lm_head...")
 
 old_model_path = "." 
-new_tokenizer_path = "./jirack_code_tokenimport torch"
-from JiRackTernary_new import JiRackConfig, JiRackTernary1B
-from transformers import AutoTokenizer
-import os
-
-print("🚀 Copying embeddings and lm_head...")
-
-old_model_path = "." ## curent model 
-new_tokenizer_path = "./jirack_code_tokenizer"  ## new JiRack tokenizer
-save_path = "./JiRack_init_model_with_new_vocab"  ## new safe bank model
+new_tokenizer_path = "./jirack_code_tokenizer" 
+save_path = "./JiRack_init_model_with_new_vocab"
 
 os.makedirs(save_path, exist_ok=True)
 
@@ -43,57 +37,23 @@ os.makedirs(save_path, exist_ok=True)
 tokenizer = AutoTokenizer.from_pretrained(new_tokenizer_path)
 new_vocab_size = len(tokenizer)
 
-# Create new model with updated vocab size
-config = JiRackConfig()
-model = JiRackTernary1B(config)
-
-# Load old model weights
-old_state = torch.load(f"{old_model_path}/pytorch_model.bin", map_location="cpu")
-
-old_vocab_size = 128256
-
-with torch.no_grad():
-    # Copy old weights
-    model.token_emb.weight[:old_vocab_size] = old_state['token_emb.weight'][:old_vocab_size]
-    model.lm_head.weight[:old_vocab_size] = old_state['lm_head.weight'][:old_vocab_size]
-    
-    # Initialize the new 3 tokens (FIM markers) with mean value
-    mean_emb = old_state['token_emb.weight'].mean(dim=0)
-    model.token_emb.weight[old_vocab_size:] = mean_emb
-    model.lm_head.weight[old_vocab_size:] = mean_emb
-
-print(f"✅ Copied {old_vocab_size} tokens")
-print(f"✅ Initialized {new_vocab_size - old_vocab_size} new tokens")
-
-# Save
-torch.save(model.state_dict(), f"{save_path}/pytorch_model.bin")
-tokenizer.save_pretrained(save_path)
-
-print(f"\n🎉 Done! New model saved to: {save_path}")
-print("Use this folder as the starting weights for training from scratch.")izer_fixed"
-save_path = "./JiRack_init_new_vocab"
-
-os.makedirs(save_path, exist_ok=True)
-
-# Load new tokenizer
-tokenizer = AutoTokenizer.from_pretrained(new_tokenizer_path)
-new_vocab_size = len(tokenizer)
+print(f"New vocab size: {new_vocab_size}")
 
-# Create new model with updated vocab size
+# Create new model
 config = JiRackConfig()
 model = JiRackTernary1B(config)
 
-# Load old model weights
-old_state = torch.load(f"{old_model_path}/pytorch_model.bin", map_location="cpu")
+# Load old model (SafeTensors)
+old_state = load_file(f"{old_model_path}/model.safetensors")
 
 old_vocab_size = 128256
 
 with torch.no_grad():
     # Copy old weights
-    model.token_emb.weight[:old_vocab_size] = old_state['token_emb.weight'][:old_vocab_size]
-    model.lm_head.weight[:old_vocab_size] = old_state['lm_head.weight'][:old_vocab_size]
+    model.token_emb.weight[:old_vocab_size] = old_state['token_emb.weight'][:old_vocab_size].clone()
+    model.lm_head.weight[:old_vocab_size] = old_state['lm_head.weight'][:old_vocab_size].clone()
     
-    # Initialize the new 3 tokens (FIM markers) with mean value
+    # Initialize new 3 tokens (FIM) with mean value
     mean_emb = old_state['token_emb.weight'].mean(dim=0)
     model.token_emb.weight[old_vocab_size:] = mean_emb
     model.lm_head.weight[old_vocab_size:] = mean_emb
@@ -101,8 +61,8 @@ with torch.no_grad():
 print(f"✅ Copied {old_vocab_size} tokens")
 print(f"✅ Initialized {new_vocab_size - old_vocab_size} new tokens")
 
-# Save
-torch.save(model.state_dict(), f"{save_path}/pytorch_model.bin")
+# Save in SafeTensors
+save_file(model.state_dict(), f"{save_path}/model.safetensors")
 tokenizer.save_pretrained(save_path)
 
 print(f"\n🎉 Done! New model saved to: {save_path}")