from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from huggingface_hub import hf_hub_download
import torch

repo = "yibinlei/effir-mistral-drop-16-mlp"

tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
config = AutoConfig.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    repo,
    config=config,
    trust_remote_code=True,
    attn_implementation="eager",
    torch_dtype=torch.bfloat16,
)

input_emb_path = hf_hub_download(repo, "embedding/input_emb.pth")
model.set_input_embeddings(torch.load(input_emb_path, map_location="cpu"))
model = PeftModel.from_pretrained(model, repo)
model = model.merge_and_unload()