Update README.md

README.md

@@ -19,11 +19,14 @@ import transformers
 import torch
 import os
 from huggingface_hub import create_repo, upload_folder
+import accelerate
 
 source_model_id = 'meta-llama/Meta-Llama-3-8B-Instruct'
 save_path = '/tmp/yujiepan/meta-llama-3-tiny-random'
 repo_id = 'yujiepan/meta-llama-3-tiny-random'
 
+os.system(f'rm -rf {save_path}')
+
 config = transformers.AutoConfig.from_pretrained(
     source_model_id,
     trust_remote_code=True,
@@ -40,6 +43,10 @@ model = transformers.AutoModelForCausalLM.from_config(
     config,
     trust_remote_code=True,
 )
+
+with accelerate.init_empty_weights():
+    model.generation_config = transformers.AutoModelForCausalLM.from_pretrained(source_model_id).generation_config
+
 model = model.to(torch.bfloat16)
 model.save_pretrained(save_path)
 
@@ -52,6 +59,8 @@ tokenizer.save_pretrained(save_path)
 model.float().generate(torch.tensor([[1, 2, 3]]).long(), max_length=16)
 
 os.system(f'ls -alh {save_path}')
+# os.system(f'rm -rf {save_path}/model.safetensors')
 create_repo(repo_id, exist_ok=True)
-upload_folder(repo_id=repo_id, folder_path=save_path)
+upload_folder(repo_id='yujiepan/meta-llama-3-tiny-random', folder_path=save_path)
+upload_folder(repo_id='yujiepan/llama-3-tiny-random', folder_path=save_path)
 ```