load_model_from_calib.py

from accelerate import Accelerator 
import logging
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

logger = logging.getLogger(__name__)

def merge_calibrators_to_hf_model(hf_model, new_tokens_start, new_tokens_end=None,
                                  embedding_calibrator=None, lm_head_calibrator=None,
                                  existing_tokens_to_calibrate=None,
                                  existing_tokens_embedding_calibrator=None, existing_tokens_lm_head_calibrator=None):
    if embedding_calibrator is not None:
        embedding_calibrator.to(hf_model.device)
    if lm_head_calibrator is not None:
        lm_head_calibrator.to(hf_model.device)
    if existing_tokens_embedding_calibrator is not None:
        existing_tokens_embedding_calibrator.to(hf_model.device)
    if existing_tokens_lm_head_calibrator is not None:
        existing_tokens_lm_head_calibrator.to(hf_model.device)

    # Handle input embeddings
    if embedding_calibrator is not None or existing_tokens_embedding_calibrator is not None:
        embedding_weights = hf_model.get_input_embeddings().weight
        with torch.no_grad():
            if embedding_calibrator is not None:
                calibrated_weights = embedding_calibrator(embedding_weights[new_tokens_start:new_tokens_end])
                try:
                    hf_model.model.embed_tokens.weight.data[new_tokens_start:new_tokens_end] = calibrated_weights
                except AttributeError:
                    # For multimodal models
                    hf_model.language_model.embed_tokens.weight.data[new_tokens_start:new_tokens_end] = calibrated_weights

            if existing_tokens_embedding_calibrator is not None and existing_tokens_to_calibrate is not None:
                existing_weights = embedding_weights[existing_tokens_to_calibrate]
                calibrated_existing = existing_tokens_embedding_calibrator(existing_weights)
                try:
                    hf_model.model.embed_tokens.weight.data[existing_tokens_to_calibrate] = calibrated_existing
                except AttributeError:
                    # For multimodal models
                    hf_model.language_model.embed_tokens.weight.data[existing_tokens_to_calibrate] = calibrated_existing

    # Handle LM head
    if lm_head_calibrator is not None or existing_tokens_lm_head_calibrator is not None:
        lm_head_weights = hf_model.get_output_embeddings().weight
        with torch.no_grad():
            if lm_head_calibrator is not None:
                calibrated_weights = lm_head_calibrator(lm_head_weights[new_tokens_start:new_tokens_end])
                hf_model.lm_head.weight.data[new_tokens_start:new_tokens_end] = calibrated_weights

            if existing_tokens_lm_head_calibrator is not None and existing_tokens_to_calibrate is not None:
                existing_weights = lm_head_weights[existing_tokens_to_calibrate]
                calibrated_existing = existing_tokens_lm_head_calibrator(existing_weights)
                hf_model.lm_head.weight.data[existing_tokens_to_calibrate] = calibrated_existing

    return hf_model



def load_model_tokenizer(
        model_name_or_path, 
        tokenizer_path=None, tokenizer=None, 
        intialization_path=None, 
        embedding_calibrator_path=None, 
        lm_calibrator_path=None, 
        num_new_words=None):
    """
    Loads a Hugging Face model and tokenizer, applies optional initialization and calibration.
    Args:
        model_name_or_path (str or AutoModelForCausalLM): The name or path of the model to load, or an already loaded model instance.
        tokenizer_path (str, optional): The name or path of the tokenizer to load. If None, the tokenizer will be loaded from the model_name_or_path.
        tokenizer (AutoTokenizer, optional): An already loaded tokenizer instance. If provided, this will be used instead of loading from tokenizer_path or model_name_or_path.
        intialization_path (str, optional): Path to the directory containing the input and output embedding state dicts for initialization. The directory should contain "input_embeddings.pt" and "output_embeddings.pt".
        embedding_calibrator_path (str, optional): Path to the embedding calibrator state dict to apply to the model.
        lm_calibrator_path (str, optional): Path to the LM head calibrator state dict to apply to the model.
        num_new_words (int, optional): The number of new words added to the tokenizer and model. This is required if either embedding_calibrator_path or lm_calibrator_path is provided, as it is needed to determine the indices of the new tokens in the model's embeddings.
    """
    if type(model_name_or_path) == str:
        logger.info("Loading model...")
        mixed_precision = "bf16" if torch.cuda.is_bf16_supported() else "fp16"
        # Load the model that we want to expand the vocabulary of
        accelerator = Accelerator(mixed_precision=mixed_precision)
    
        model = AutoModelForCausalLM.from_pretrained(
            model_name_or_path, 
            torch_dtype=torch.bfloat16 if mixed_precision == "bf16" else torch.float16
            )
        model = accelerator.prepare(model)
        model.eval()
        if tokenizer_path is not None:
            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
        else:
            tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
        

    elif type(model_name_or_path) == AutoModelForCausalLM:
        model = model_name_or_path
        tokenizer = tokenizer
    else:
        raise ValueError("model_name_or_path must be a string or an AutoModelForCausalLM instance.")
    
    # Load initialization if path provided
    if intialization_path is not None:
        logger.info("Loading initialization...")
        inp_emb_file_path = intialization_path + "/input_embeddings.pt"
        lm_head_emb_file_path = intialization_path + "/output_embeddings.pt"
        input_embeddings_state_dict = torch.load(inp_emb_file_path,map_location=torch.device('cpu'))
        lm_head_embeddings_state_dict = torch.load(lm_head_emb_file_path,map_location=torch.device('cpu'))

        # Resize embedding and lm head to match the loaded state dict
        model.get_input_embeddings().weight = torch.nn.Parameter(torch.zeros_like(input_embeddings_state_dict['weight']))
        model.get_output_embeddings().weight = torch.nn.Parameter(torch.zeros_like(lm_head_embeddings_state_dict['weight']))
        model.get_input_embeddings().load_state_dict(input_embeddings_state_dict)
        model.get_output_embeddings().load_state_dict(lm_head_embeddings_state_dict)
        # Ensure that the weights are on the same device using accelerator
        model = accelerator.prepare(model)


    if embedding_calibrator_path is not None:
        logger.info("Applying embedding calibrator...")
        embedding_calibrator = torch.load(embedding_calibrator_path, weights_only=False, map_location=torch.device('cpu'))
    if lm_calibrator_path is not None:
        logger.info("Applying LM calibrator...")
        lm_calibrator = torch.load(lm_calibrator_path, weights_only=False, map_location=torch.device('cpu'))
    
    if embedding_calibrator_path is not None or lm_calibrator_path is not None:
        logger.info("Merging calibrators into model...")
        assert num_new_words is not None, "num_new_words must be provided when applying calibrators."
        model = merge_calibrators_to_hf_model(
            model, 
            new_tokens_start= int(model.get_input_embeddings().weight.shape[0])-num_new_words,
            new_tokens_end= int(model.get_input_embeddings().weight.shape[0]),
            embedding_calibrator=embedding_calibrator if embedding_calibrator_path is not None else None,
            lm_head_calibrator=lm_calibrator if lm_calibrator_path is not None else None
            )


    return model, tokenizer

if __name__ == "__main__":
    model_name_or_path = "Qwen/Qwen3-30B-A3B-Instruct-2507"
    
    FOLDER_PATH = "" # e.g. Path leading to the experiment folder named some thing like "gujarati_qwen3_30b_glot500c_1000"
    tokenizer_path =f"{FOLDER_PATH}/models/gujarati_qwen3_30b_glot500c_1000_tokenizer.pt"
    intialization_path = f"{FOLDER_PATH}/embeddings/"
    embedding_calibrator_path = f"{FOLDER_PATH}/calibrators/embedding_calibrator.pt"
    lm_calibrator_path = f"{FOLDER_PATH}/calibrators/lm_head_calibrator.pt"
    metrics_file_path = f"{FOLDER_PATH}/metrics_other.json"

    with open(metrics_file_path, "r") as f:
        import json
        num_new_words = json.load(f)["n_new_words"]["expanded"]
   
    model, tokenizer = load_model_tokenizer(
        model_name_or_path=model_name_or_path,
        tokenizer_path=tokenizer_path,
        intialization_path=intialization_path,
        embedding_calibrator_path=embedding_calibrator_path,
        lm_calibrator_path=lm_calibrator_path,
        num_new_words=num_new_words
    )