Refactoring: more concise and readable

convert-pth-to-ggml.py

@@ -17,6 +17,7 @@
 # and vocabulary.
 #
 
+import argparse
 import sys
 import json
 import struct
@@ -24,45 +25,29 @@ import numpy as np
 import torch
 from sentencepiece import SentencePieceProcessor
 
-if len(sys.argv) < 3:
+parser = argparse.ArgumentParser(description='Convert a LLaMA model checkpoint to a ggml compatible file')
-    print("Usage: convert-ckpt-to-ggml.py dir-model ftype\n")
+parser.add_argument('dir_model', help='directory containing the model checkpoint')
-    print("  ftype == 0 -> float32")
+parser.add_argument('ftype', type=int, choices=[0, 1], default=1, help='file type (0: float32, 1: float16)')
-    print("  ftype == 1 -> float16")
+args = parser.parse_args()
-    sys.exit(1)
 
 # output in the same directory as the model
-dir_model = sys.argv[1]
+dir_model = args.dir_model
+ftype = args.ftype
 
 fname_hparams   = f"{dir_model}/params.json"
 fname_tokenizer = f"{dir_model}/../tokenizer.model"
 
 def get_n_parts(dim):
-    mappings = {
+    mappings = {4096: 1, 5120: 2, 6656: 4, 8192: 8}
-        4096: 1,
+    n_parts = mappings.get(dim)
-        5120: 2,
+    if n_parts is None:
-        6656: 4,
-        8192: 8
-    }
-    if dim not in mappings:
         print(f"Invalid dim: {dim}")
         sys.exit(1)
-    return mappings[dim]
+    return n_parts
 
-# possible data types
-#   ftype == 0 -> float32
-#   ftype == 1 -> float16
-#
 # map from ftype to string
 ftype_str = ["f32", "f16"]
 
-ftype = 1
-if len(sys.argv) > 2:
-    ftype = int(sys.argv[2])
-    if ftype < 0 or ftype > 1:
-        print(f"Invalid ftype: {ftype}")
-        sys.exit(1)
-    fname_out = f"{dir_model}/ggml-model-{ftype_str[ftype]}.bin"
-
 with open(fname_hparams, "r") as f:
     hparams = json.load(f)
 
@@ -78,15 +63,13 @@ print(f"n_parts = {n_parts}\n")
 for p in range(n_parts):
     print(f"Processing part {p}\n")
 
-    #fname_model = sys.argv[1] + "/consolidated.00.pth"
+    #fname_model = args[1] + "/consolidated.00.pth"
     fname_model = f"{dir_model}/consolidated.0{p}.pth"
-    fname_out = f"{dir_model}/ggml-model-{ftype_str[ftype]}.bin"
+    fname_out = f"{dir_model}/ggml-model-{ftype_str[ftype]}.bin{'' if p == 0 else '.' + str(p)}"
-    if (p > 0):
-        fname_out = f"{dir_model}/ggml-model-{ftype_str[ftype]}.bin.{p}"
 
     model = torch.load(fname_model, map_location="cpu")
 
-    fout = open(fname_out, "wb")
+    with open(fname_out, "wb") as fout:
 
         fout.write(struct.pack("i", 0x67676d6c)) # magic: ggml in hex
         fout.write(struct.pack("i", hparams["vocab_size"]))
@@ -122,19 +105,18 @@ for p in range(n_parts):
                 fout.write(struct.pack("i", len(text)))
                 fout.write(text)
 
-    for k, v in model.items():
+        for name, data in model.items():
-        name = k
+            shape = data.shape
-        shape = v.shape
 
             # skip layers.X.attention.inner_attention.rope.freqs
-        if name[-5:] == "freqs":
+            if name.endswith("freqs"):
                 continue
 
-        print(f"Processing variable: {name} with shape: {data.shape} and type: {data.dtype}\n")
+            print(f"Processing variable: {name} with shape: {shape} and type: {data.dtype}\n")
 
             #data = tf.train.load_variable(dir_model, name).squeeze()
-        data = v.numpy().squeeze()
+            data = np.squeeze(data)
-        n_dims = len(data.shape);
+            n_dims = len(data.shape)
 
             # for efficiency - transpose some matrices
             # "model/h.*/attn/c_attn/w"
@@ -148,8 +130,6 @@ for p in range(n_parts):
             #    print("  Transposing")
             #    data = data.transpose()
 
-        dshape = data.shape
-
             # default type is fp16
             ftype_cur = 1
             if ftype == 0 or n_dims == 1:
@@ -160,16 +140,14 @@ for p in range(n_parts):
             # header
             sname = name.encode('utf-8')
             fout.write(struct.pack("iii", n_dims, len(sname), ftype_cur))
-        for i in range(n_dims):
+            for dim in reversed(data.shape):
-            fout.write(struct.pack("i", dshape[n_dims - 1 - i]))
+                fout.write(struct.pack("i", dim))
-        fout.write(sname);
+
+            fout.write(sname)
 
             # data
             data.tofile(fout)
 
-    # I hope this deallocates the memory ..
+        del model
-    model = None
-
-    fout.close()
 
     print(f"Done. Output file: {fname_out}, (part {p})\n")