clean

ggml-org · diimdeep · Mar 15, 2023 · Mar 15, 2023 · Mar 15, 2023 · Mar 15, 2023
commit 0deb075a3f8bf5445f4148a2d75c8a179d42da5b
diff --git a/convert-pth-to-ggml.py b/convert-pth-to-ggml.py
@@ -84,6 +84,50 @@ def get_n_parts(dim):
 print(f'Parts to process: {n_parts}')
 
 
+def load_model(fname):
+    class Tensor():
+        def __init__(self, shape, dtype, loadinfo):
+            self.shape = shape
+            self.dtype = dtype
+            self.loadinfo = loadinfo
+
+        def numpy(self):
+            myzip, base_name, storage_offset, k, shape, dtype = self.loadinfo
+            with myzip.open(f'{base_name}/data/{k}') as myfile:
+                bytes_size = np.dtype(self.dtype).itemsize
+                myfile.seek(storage_offset * bytes_size, 1)
+                ret = np.empty(shape, dtype=dtype)
+                myfile.readinto(ret.data)
+                return ret
+
+    def my_unpickle(datapkl, myzip, base_name):
+        def my_rebuild_tensor(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata=None):
+            storage_type = storage[1]
+            obj_key = storage[2]
+            return Tensor(shape=size, dtype=storage_type, loadinfo=(
+                myzip, base_name, storage_offset,
+                obj_key, size, storage_type
+            ))
+
+        class MyUnpickler(pickle.Unpickler):
+            def find_class(self, *p):
+                if p == ('torch', 'HalfStorage'): return np.float16
+                if p == ('torch', 'FloatStorage'): return np.float32
+                if p == ('torch._utils', '_rebuild_tensor_v2'): return my_rebuild_tensor
+                if p == ('collections', 'OrderedDict'): return dict
+                raise ValueError(f'Unrecognized pickle {p}')
+
+            def persistent_load(self, pid):
+                return pid
+
+        return MyUnpickler(datapkl).load()
+
+    myzip =  zipfile.ZipFile(fname, 'r')
+    base_name = myzip.namelist()[0].split('/', 1)[0]
+    with myzip.open(f'{base_name}/data.pkl') as myfile:
+        model = my_unpickle(myfile, myzip, base_name)
+    return model
+
 def get_fname(p):
     fname = "/consolidated.0" + str(p) + ".pth"
     return fname
@@ -133,51 +177,6 @@ def process_part(p):
             fout.write(struct.pack("i", len(text)))
             fout.write(text)
 
-
-    def load_model(fname):
-        class Tensor():
-            def __init__(self, shape, dtype, loadinfo):
-                self.shape = shape
-                self.dtype = dtype
-                self.loadinfo = loadinfo
-
-            def numpy(self):
-                myzip, base_name, storage_offset, k, shape, dtype = self.loadinfo
-                with myzip.open(f'{base_name}/data/{k}') as myfile:
-                    bytes_size = np.dtype(self.dtype).itemsize
-                    myfile.seek(storage_offset * bytes_size, 1)
-                    ret = np.empty(shape, dtype=dtype)
-                    myfile.readinto(ret.data)
-                    return ret
-
-        def my_unpickle(datapkl, myzip, base_name):
-            def my_rebuild_tensor(storage, storage_offset, size, stride, requires_grad, backward_hooks, metadata=None):
-                storage_type = storage[1]
-                obj_key = storage[2]
-                return Tensor(shape=size, dtype=storage_type, loadinfo=(
-                    myzip, base_name, storage_offset,
-                    obj_key, size, storage_type
-                ))
-
-            class MyUnpickler(pickle.Unpickler):
-                def find_class(self, *p):
-                    if p == ('torch', 'HalfStorage'): return np.float16
-                    if p == ('torch', 'FloatStorage'): return np.float32
-                    if p == ('torch._utils', '_rebuild_tensor_v2'): return my_rebuild_tensor
-                    if p == ('collections', 'OrderedDict'): return dict
-                    raise ValueError(f'Unrecognized pickle {p}')
-
-                def persistent_load(self, pid):
-                    return pid
-
-            return MyUnpickler(datapkl).load()
-
-        myzip =  zipfile.ZipFile(fname, 'r')
-        base_name = myzip.namelist()[0].split('/', 1)[0]
-        with myzip.open(f'{base_name}/data.pkl') as myfile:
-            model = my_unpickle(myfile, myzip, base_name)
-        return model
-
     model = load_model(fname_model)
 
     q = queue.Queue(maxsize=2)
@@ -234,7 +233,6 @@ def writer():
 
         # data
         memout.write(data.tobytes())
-        # data.tofile(memout)
         q.put(memout)
 
     q.join()