Function calibrate_model

examples/mediatek/model_export_scripts/gemma.py:352–376

(model, cal_dataset, chunk_idx: str)

def calibrate_model(model, cal_dataset, chunk_idx: str):
    with torch.no_grad():
        for inp in tqdm(cal_dataset, desc="Calibrating Model: "):
            # pass prompt and response
            for batch in tqdm(inp[chunk_idx].keys(), desc="Batch: "):
                if inp[chunk_idx][batch] is not None:
                    inputs_embeds = torch.tensor(inp[chunk_idx][batch]["hidden_state"])
                    pos_emb = torch.tensor(inp[chunk_idx][batch]["pos_emb"])
                    cache = torch.tensor(inp[chunk_idx][batch]["cache"])
                    mask = inp[chunk_idx][batch]["mask"]
                    if isinstance(mask, dict):
                        global_mask = torch.tensor(mask["GLOBAL"])
                        local_mask = torch.tensor(mask["SLIDING_LOCAL"])
                        model(
                            inputs_embeds,
                            global_mask,
                            local_mask,
                            pos_emb,
                            *torch.split(cache, 1, dim=0),
                        )
                    else:
                        mask = torch.tensor(mask)
                        model(
                            inputs_embeds, mask, pos_emb, *torch.split(cache, 1, dim=0)
                        )


def export_to_et_ir(