(model, cal_dataset, chunk_idx: str)
| 310 | |
| 311 | |
| 312 | def calibrate_model(model, cal_dataset, chunk_idx: str): |
| 313 | with torch.no_grad(): |
| 314 | for inp in tqdm(cal_dataset, desc="Calibrating Model: "): |
| 315 | # pass prompt and response |
| 316 | for batch in tqdm(inp[chunk_idx].keys(), desc="Batch: "): |
| 317 | if inp[chunk_idx][batch] is not None: |
| 318 | inputs_embeds = torch.tensor(inp[chunk_idx][batch]["hidden_state"]) |
| 319 | mask = torch.tensor(inp[chunk_idx][batch]["mask"]) |
| 320 | pos_emb = torch.tensor(inp[chunk_idx][batch]["pos_emb"]) |
| 321 | cache = torch.tensor(inp[chunk_idx][batch]["cache"]) |
| 322 | model(inputs_embeds, mask, pos_emb, *torch.split(cache, 1, dim=0)) |
| 323 | |
| 324 | |
| 325 | def export_to_et_ir( |
no test coverage detected