Attempt to compress code

Author: metric-space

mergekit/merge.py

@@ -15,7 +15,7 @@
 
 import logging
 import os
-from typing import Optional
+from typing import Dict, List, Optional, Tuple
 
 import tqdm
 import transformers
@@ -26,6 +26,7 @@
     get_architecture_info,
 )
 from mergekit.card import generate_card
+from mergekit.common import ModelReference
 from mergekit.config import MergeConfiguration
 from mergekit.graph import Executor
 from mergekit.io.tasks import LoaderCache
@@ -46,31 +47,27 @@ def run_merge(
     if not merge_config.models and not merge_config.slices:
         raise RuntimeError("No output requested")
 
-    ## TODO: ----------- reconcile these steps -------------
-
-    model_arch_info = [
-        get_architecture_info(m.config(trust_remote_code=options.trust_remote_code))
-        for m in merge_config.referenced_models()
-    ]
-
-    arch_dict = {
-        m.model.path: ConfiguredArchitectureInfo(
-            info=get_architecture_info(
-                m.config(trust_remote_code=options.trust_remote_code)
+    model_arch_info: List[Tuple[ModelReference, ConfiguredArchitectureInfo]] = [
+        (
+            m,
+            ConfiguredArchitectureInfo(
+                info=get_architecture_info(
+                    m.config(trust_remote_code=options.trust_remote_code)
+                ),
+                config=m.config(trust_remote_code=options.trust_remote_code),
             ),
-            config=m.config(trust_remote_code=options.trust_remote_code),
         )
         for m in merge_config.referenced_models()
-    }
-
-    ## ----------------------------------------------------
+    ]
 
     if not options.allow_crimes:
-        if not all(a == model_arch_info[0] for a in model_arch_info[1:]):
+        if not all(
+            a[1].info == model_arch_info[0][1].info for a in model_arch_info[1:]
+        ):
             raise RuntimeError(
                 "Must specify --allow-crimes to attempt to mix different architectures"
             )
-    arch_info = model_arch_info[0]
+    arch_info: ArchitectureInfo = model_arch_info[0][1].info
 
     # initialize loader cache and set options
     loader_cache = LoaderCache()
@@ -88,7 +85,7 @@ def run_merge(
     targets = MergePlanner(
         merge_config,
         arch_info,
-        arch_dict,
+        dict(model_arch_info),
         out_path=out_path,
         options=options,
         out_model_config=cfg_out,

mergekit/plan.py

@@ -44,7 +44,7 @@
 class MergePlanner:
     config: MergeConfiguration
     arch_info: ArchitectureInfo
-    arch_dict: Dict[str, ConfiguredArchitectureInfo]
+    arch_dict: Dict[ModelReference, ConfiguredArchitectureInfo]
     clone_tensors: bool
     trust_remote_code: bool
     out_model_config: Any
@@ -60,7 +60,7 @@ def __init__(
         config: MergeConfiguration,
         arch_info: ArchitectureInfo,
         arch_dict: Dict[
-            str, ConfiguredArchitectureInfo
+            ModelReference, ConfiguredArchitectureInfo
         ],  # perhaps this should no longer be a disjoint step
         out_path: str,
         options: MergeOptions,
@@ -216,9 +216,7 @@ def plan_layer(
             config=self.out_model_config,
         )
         weights_in: List[List[WeightInfo]] = [
-            self.arch_dict[s.model.model.path].layer_weights(
-                index=s.layer_range[0] + layer_offset
-            )
+            self.arch_dict[s.model].layer_weights(index=s.layer_range[0] + layer_offset)
             for s in sources
         ]
 
@@ -268,9 +266,7 @@ def plan(self):
         models_ = [s.model for s in self.config.slices[0].sources]
         for weight_infos in zip(
             *[
-                self.arch_dict[m.model.path].info.pre_weights(
-                    config=self.out_model_config
-                )
+                self.arch_dict[m].info.pre_weights(config=self.out_model_config)
                 for m in models_
             ]
         ):
@@ -291,9 +287,7 @@ def plan(self):
         models_ = [s.model for s in self.config.slices[-1].sources]
         for weight_infos in zip(
             *[
-                self.arch_dict[m.model.path].info.post_weights(
-                    config=self.out_model_config
-                )
+                self.arch_dict[m].info.post_weights(config=self.out_model_config)
                 for m in models_
             ]
         ):