import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
from tqdm import tqdm

# Set environment variables and configure the number of threads
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
torch.set_num_threads(1)

# Define model names and output directory
base_model_name1 = "huihui-ai/QwQ-32B-Preview-abliterated"
base_model_name2 = "huihui-ai/Qwen2.5-Coder-32B-Instruct-abliterated"
output_model_dir = "huihui-ai/QwQ-32B-Coder-Fusion-9010"  # Directory to save the merged model

# Specify the device for computation
device = torch.device("cpu")

# Load models and tokenizer onto the CPU
base_model1 = AutoModelForCausalLM.from_pretrained(base_model_name1, torch_dtype=torch.bfloat16).to(device)
base_model2 = AutoModelForCausalLM.from_pretrained(base_model_name2, torch_dtype=torch.bfloat16).to(device)
tokenizer = AutoTokenizer.from_pretrained(base_model_name1)

def merge_model_weights(base_model1, base_model2, alpha=0.9):
    """
    Merge the weights of two models based on a specified ratio and return the updated model.
    The parameter alpha determines the blending ratio, with a default of 0.9 
    (90% from base_model1 and 10% from base_model2).
    """
    # Merge weights of the output layers
    base_model1.lm_head.weight.data = (alpha * base_model1.lm_head.weight.data + (1 - alpha) * base_model2.lm_head.weight.data).to(device)
    base_model1.model.embed_tokens.weight.data = (alpha * base_model1.model.embed_tokens.weight.data + (1 - alpha) * base_model2.model.embed_tokens.weight.data).to(device)
    
    # Merge weights for each transformer layer
    with tqdm(total=len(base_model1.model.layers), desc="Merging weights for layers") as pbar:
        for i in range(len(base_model1.model.layers)):
            base_model1.model.layers[i].input_layernorm.weight.data = (alpha * base_model1.model.layers[i].input_layernorm.weight.data + (1 - alpha) * base_model2.model.layers[i].input_layernorm.weight.data).to(device)
            base_model1.model.layers[i].mlp.down_proj.weight.data = (alpha * base_model1.model.layers[i].mlp.down_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].mlp.down_proj.weight.data).to(device)
            base_model1.model.layers[i].mlp.gate_proj.weight.data = (alpha * base_model1.model.layers[i].mlp.gate_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].mlp.gate_proj.weight.data).to(device)
            base_model1.model.layers[i].mlp.up_proj.weight.data = (alpha * base_model1.model.layers[i].mlp.up_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].mlp.up_proj.weight.data).to(device)
            base_model1.model.layers[i].post_attention_layernorm.weight.data = (alpha * base_model1.model.layers[i].post_attention_layernorm.weight.data + (1 - alpha) * base_model2.model.layers[i].post_attention_layernorm.weight.data).to(device)
            base_model1.model.layers[i].self_attn.q_proj.weight.data = (alpha * base_model1.model.layers[i].self_attn.q_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].self_attn.q_proj.weight.data).to(device)
            base_model1.model.layers[i].self_attn.k_proj.weight.data = (alpha * base_model1.model.layers[i].self_attn.k_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].self_attn.k_proj.weight.data).to(device)
            base_model1.model.layers[i].self_attn.v_proj.weight.data = (alpha * base_model1.model.layers[i].self_attn.v_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].self_attn.v_proj.weight.data).to(device)
            base_model1.model.layers[i].self_attn.o_proj.weight.data = (alpha * base_model1.model.layers[i].self_attn.o_proj.weight.data + (1 - alpha) * base_model2.model.layers[i].self_attn.o_proj.weight.data).to(device)
            
            pbar.update(1)  # Update the progress bar

    # Merge weights for the final normalization layer
    base_model1.model.norm.weight.data = (alpha * base_model1.model.norm.weight.data + (1 - alpha) * base_model2.model.norm.weight.data).to(device)
    
    return base_model1

# Merge the weights of the two models with a blending ratio of 0.9
merged_model = merge_model_weights(base_model1, base_model2, alpha=0.9)

# Save the merged model and tokenizer
merged_model.save_pretrained(output_model_dir)
tokenizer.save_pretrained(output_model_dir)
print(f"Merged model and tokenizer saved to {output_model_dir}")