Spaces:
Sleeping
Sleeping
File size: 4,103 Bytes
b5e9acb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 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 |