From 530ef8e778382b23ba5472b6c796203fde7053cb Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 27 Jul 2023 05:08:45 +0000 Subject: [PATCH] light touchups to export script so one doesn't need to pass in a slash at the end --- export_meta_llama_bin.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/export_meta_llama_bin.py b/export_meta_llama_bin.py index 801077b..3d07c1c 100644 --- a/export_meta_llama_bin.py +++ b/export_meta_llama_bin.py @@ -1,6 +1,7 @@ """ This script exports the Llama 2 weights in llama2c.bin format. """ +import os import sys import struct from pathlib import Path @@ -89,16 +90,13 @@ def concat_weights(models): def load_and_export(model_path, output_path): - with open(model_path + 'params.json') as f: + params_path = os.path.join(model_path, 'params.json') + with open(params_path) as f: params = json.load(f) print(params) model_paths = sorted(list(Path(model_path).glob('consolidated.*.pth'))) - models = [] - for i in model_paths: - print(f'Loading {i}') - models.append(torch.load(i, map_location='cpu')) - + models = [torch.load(p, map_location='cpu') for p in model_paths] state_dict = concat_weights(models) del models export(params, state_dict, output_path)