import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model

def main(seed):
    torch.manual_seed(seed)
    model = AutoModelForCausalLM.from_pretrained(
        "facebook/opt-125m",
        quantization_config=BitsAndBytesConfig(load_in_8bit=True),
    )
    random_input = torch.LongTensor([[1, 0, 1, 0, 1, 0]]).to(model.device)
    out_base = F.softmax(model(random_input).logits, dim=-1)
    config = LoraConfig(
        r=8,
        init_lora_weights=False,
        lora_bias=True,
    )
    model = get_peft_model(model, config)
    with torch.inference_mode():
        out_before_merge = F.softmax(model(random_input).logits, dim=-1)

    model.merge_and_unload()
    with torch.inference_mode():
        out_after_merge = F.softmax(model(random_input).logits, dim=-1)

    atol = 1e-3
    rtol = 1
    if torch.allclose(out_before_merge, out_after_merge, atol=atol, rtol=rtol):
        seed_pass.append(seed)
        print(f"seed: {seed} " + "pass " + "="*10)
    else:
        seed_fail.append(seed)
        print(f"seed: {seed} " + "fail " + "e"*20)

if __name__ == "__main__":
    seed_pass = []
    seed_fail = []
    for seed in range(100):
        main(seed)
    print(f"passed seeds: {seed_pass}")
    print(f"failed seeds: {seed_fail}")