Create merge_utils.py

merge_utils.py

@@ -0,0 +1,128 @@
+import os
+import yaml
+import gc
+import torch
+import shutil
+from pathlib import Path
+
+# --- CRITICAL PATCH: MUST RUN BEFORE MERGEKIT IMPORTS ---
+import pydantic
+from pydantic import ConfigDict, BaseModel
+BaseModel.model_config = ConfigDict(arbitrary_types_allowed=True)
+
+try:
+    from mergekit.config import MergeConfiguration
+    from mergekit.merge import run_merge
+    from mergekit.architecture import get_architecture_info
+except ImportError:
+    print("MergeKit not found. Please install 'mergekit' in requirements.txt")
+
+def execute_mergekit_config(config_dict, out_path, shard_gb, device="cpu"):
+    """
+    Executes a MergeKit run based on a dictionary config.
+    Optimized for CPU execution with aggressive sharding.
+    """
+    # Convert dict to YAML string first to ensure validation passes through standard flow
+    config_yaml = yaml.dump(config_dict)
+    
+    print("--- Generated MergeKit Config ---")
+    print(config_yaml)
+    print("---------------------------------")
+
+    conf = MergeConfiguration.model_validate(yaml.safe_load(config_yaml))
+
+    run_merge(
+        conf,
+        out_path=out_path,
+        device=device,
+        low_cpu_mem=True,
+        copy_tokenizer=True,
+        lazy_unpickle=True,
+        max_shard_size=int(shard_gb * 1024**3)
+    )
+    
+    # Force cleanup
+    gc.collect()
+
+def build_full_merge_config(
+    method, models, base_model, weights, density, 
+    dtype, tokenizer_source, layer_ranges
+):
+    """
+    Constructs the YAML dictionary for general merging (Linear, SLERP, TIES, etc.)
+    """
+    # Basic Config
+    config = {
+        "merge_method": method.lower(),
+        "base_model": base_model if base_model else models[0],
+        "dtype": dtype,
+        "tokenizer_source": tokenizer_source,
+        "models": []
+    }
+
+    # Helper to parse weights safely
+    w_list = []
+    if weights:
+        try:
+            w_list = [float(x.strip()) for x in weights.split(',')]
+        except:
+            print("Warning: Could not parse weights, defaulting to 1.0")
+    
+    # Model Construction
+    for i, m in enumerate(models):
+        entry = {"model": m, "parameters": {}}
+        
+        # Method Specific Param Injection
+        if method.lower() in ["ties", "dare_ties", "dare_linear"]:
+            entry["parameters"]["weight"] = w_list[i] if i < len(w_list) else 1.0
+            entry["parameters"]["density"] = density
+        
+        elif method.lower() == "slerp":
+            # SLERP usually takes 't' parameter via weight, but often requires layer slices
+            # If layer_ranges is provided (JSON), use that. Otherwise use weight as 't'
+            if layer_ranges and "slices" in layer_ranges:
+                # Advanced Slice Config
+                pass # mergekit handles slices at root level usually, but we inject here if needed
+            else:
+                entry["parameters"]["weight"] = w_list[i] if i < len(w_list) else 1.0
+
+        elif method.lower() == "linear":
+             entry["parameters"]["weight"] = w_list[i] if i < len(w_list) else 1.0
+
+        config["models"].append(entry)
+    
+    # Inject Slices/Layer Ranges if provided (Raw JSON override)
+    if layer_ranges.strip():
+        try:
+            extra_params = yaml.safe_load(layer_ranges)
+            if isinstance(extra_params, dict):
+                config.update(extra_params)
+        except Exception as e:
+            print(f"Error parsing layer ranges JSON: {e}")
+
+    return config
+
+def build_moe_config(
+    base_model, experts, gate_mode, dtype, 
+    tokenizer_source, positive_prompts=None
+):
+    """
+    Constructs the YAML dictionary for Mixture of Experts (MoE)
+    """
+    config = {
+        "base_model": base_model,
+        "gate_mode": gate_mode,
+        "dtype": dtype,
+        "tokenizer_source": tokenizer_source,
+        "experts": []
+    }
+    
+    # Parse experts
+    for i, exp in enumerate(experts):
+        expert_entry = {
+            "source_model": exp,
+            "positive_prompts": [f"expert_{i}"] # Placeholder if not provided
+        }
+        config["experts"].append(expert_entry)
+        
+    return config