DeepWater-Pleroma-12B-v0-raw-weights

import torch
from safetensors.torch import load_file, save_file
import os
import gc

# Configuration
base_path = r'B:\12B\models--mistralai--Mistral-Nemo-Instruct-2407'
broken_path = r'C:\Quanter\model_cache\EldritchLabs__DeepWater-Pleroma-12B-v1'
output_path = r'C:\Quanter\model_cache\EldritchLabs__DeepWater-Pleroma-12B-v1\DeepWater-Healed'
os.makedirs(output_path, exist_ok=True)

print("Step 1: Indexing base model shards...")
base_map = {}
base_files = [f for f in os.listdir(base_path) if f.endswith('.safetensors')]
for bf in base_files:
    # We only load the header to index keys (fast)
    from safetensors import safe_open
    with safe_open(os.path.join(base_path, bf), framework="pt") as f:
        for k in f.keys():
            base_map[k] = bf

def get_base_tensor(name):
    """Helper to find and load a specific tensor from the base model."""
    if name not in base_map:
        return None
    target_file = os.path.join(base_path, base_map[name])
    sd = load_file(target_file)
    return sd[name].clone().detach()

print("Step 2: Healing broken shards...")
broken_files = [f for f in os.listdir(broken_path) if f.endswith('.safetensors')]

for bf in broken_files:
    print(f"Processing {bf}...")
    broken_file_path = os.path.join(broken_path, bf)
    
    # Load broken shard into RAM and break disk link
    mmap_broken = load_file(broken_file_path)
    broken_sd = {k: v.clone().detach() for k, v in mmap_broken.items()}
    del mmap_broken
    gc.collect()
    
    shard_healed = False
    for k in list(broken_sd.keys()):
        broken_t = broken_sd[k]
        
        # Check for inf/nan
        invalid_mask = ~torch.isfinite(broken_t)
        if invalid_mask.any():
            num_broken = torch.sum(invalid_mask).item()
            print(f"  !! Found {num_broken} inf/nan in {k}. Fetching base weights...")
            
            base_t = get_base_tensor(k)
            if base_t is not None:
                # Ensure shapes match (handle potential mergekit resizing)
                if base_t.shape != broken_t.shape:
                    print(f"     Shape mismatch for {k}: Base {base_t.shape} vs Broken {broken_t.shape}. Skipping.")
                    continue
                
                # Heal: Keep broken where finite, take base where infinite
                broken_sd[k] = torch.where(invalid_mask, base_t, broken_t)
                shard_healed = True
                del base_t
            else:
                print(f"     Warning: {k} not found in base model. Cannot heal.")

    # Save the healed shard to the NEW directory
    save_file(broken_sd, os.path.join(output_path, bf))
    print(f"  Saved {bf}")
    
    del broken_sd
    gc.collect()

print("\nHealing complete. Output saved to:", output_path)