Create dare_utils.py

dare_utils.py

@@ -0,0 +1,94 @@
+import torch
+import os
+import gc
+from safetensors.torch import load_file, save_file
+from huggingface_hub import hf_hub_download, HfApi
+from tqdm import tqdm
+import numpy as np
+
+def download_file(repo_id, filename, token, local_dir):
+    return hf_hub_download(repo_id=repo_id, filename=filename, token=token, local_dir=local_dir)
+
+def task_dare_custom(hf_token, base_repo, fine_tuned_repo, ratio, mask_rate, out_repo, private):
+    """
+    Custom DARE implementation:
+    1. Load Base & Fine-Tuned
+    2. Delta = FT - Base
+    3. Mask = Bernoulli(1 - mask_rate)
+    4. Rescale = 1 / (1 - mask_rate)
+    5. New = Base + (Delta * Mask * Rescale * Ratio)
+    """
+    api = HfApi(token=hf_token)
+    temp_dir = "./temp_dare"
+    os.makedirs(temp_dir, exist_ok=True)
+
+    try:
+        # 1. Identify Model Files (Naive: get first .safetensors)
+        print("Locating files...")
+        base_files = api.list_repo_files(base_repo)
+        base_sf = next((f for f in base_files if f.endswith(".safetensors") and "model" in f), None)
+        
+        ft_files = api.list_repo_files(fine_tuned_repo)
+        ft_sf = next((f for f in ft_files if f.endswith(".safetensors") and "model" in f), None)
+
+        if not base_sf or not ft_sf:
+            return "Error: Could not locate .safetensors in one of the repos."
+
+        print(f"Downloading {base_sf}...")
+        base_path = download_file(base_repo, base_sf, hf_token, temp_dir)
+        print(f"Downloading {ft_sf}...")
+        ft_path = download_file(fine_tuned_repo, ft_sf, hf_token, temp_dir)
+
+        # 2. Process
+        print("Loading tensors...")
+        base_sd = load_file(base_path, device="cpu")
+        ft_sd = load_file(ft_path, device="cpu")
+        
+        merged_sd = {}
+        keys = set(base_sd.keys()).intersection(ft_sd.keys())
+        
+        scale_factor = 1.0 / (1.0 - mask_rate)
+        
+        print("Applying DARE...")
+        for k in tqdm(keys):
+            b_tensor = base_sd[k]
+            f_tensor = ft_sd[k]
+            
+            if b_tensor.shape != f_tensor.shape:
+                merged_sd[k] = f_tensor # Fallback
+                continue
+            
+            # Skip 1D tensors (LayerNorms usually) or non-float
+            if len(b_tensor.shape) < 2 or not b_tensor.is_floating_point():
+                merged_sd[k] = f_tensor # Keep FT version
+                continue
+
+            # Calculate Delta
+            delta = f_tensor - b_tensor
+            
+            # Create Mask (Bernoulli)
+            mask = torch.bernoulli(torch.full_like(delta, 1.0 - mask_rate))
+            
+            # Apply DARE formula
+            # New = Base + (Delta * Mask * Scale * Ratio)
+            final = b_tensor + (delta * mask * scale_factor * ratio)
+            
+            merged_sd[k] = final.to(torch.bfloat16) # Enforce BF16 for save
+
+        # 3. Save
+        out_path = os.path.join(temp_dir, "dare_merged.safetensors")
+        save_file(merged_sd, out_path)
+        
+        # 4. Upload
+        print("Uploading...")
+        api.create_repo(repo_id=out_repo, private=private, exist_ok=True)
+        api.upload_file(path_or_fileobj=out_path, path_in_repo="model.safetensors", repo_id=out_repo)
+        
+        return f"Done! Uploaded to {out_repo}"
+
+    except Exception as e:
+        return f"DARE Error: {e}"
+    finally:
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+        gc.collect()