Views
No views yet
1import torch
2from bah.models import ConflictAwareAHModel
3from huggingface_hub import hf_hub_download
4
5ckpt_path = hf_hub_download(repo_id="Bekhouche/ConflictAwareAH", filename="best_model.pt")
6ckpt = torch.load(ckpt_path, map_location="cpu")
7args = ckpt["args"]
8
9# Infer fusion_type from checkpoint keys
10state_keys = set(ckpt["model"].keys())
11fusion_type = args.get("fusion_type") or ("6token" if any("fusion_transformer" in k for k in state_keys) else "concat")
12
13model = ConflictAwareAHModel(
14 video_model=args["video_model"],
15 audio_model=args["audio_model"],
16 text_model=args["text_model"],
17 dropout=0.0,
18 freeze_encoders=args.get("freeze_encoders", True),
19 unfreeze_top_k=args.get("unfreeze_top_k", 0),
20 num_transformer_layers=args.get("num_layers", 2),
21 fusion_type=fusion_type,
22)
23model.load_state_dict(ckpt["model"], strict=True)
24model.eval()
25
26text_blend = ckpt.get("text_blend", args.get("text_blend", 0.5))