| |
| |
| |
| |
| |
| import os, sys, io, importlib.util, tempfile, numpy as np, torch, soundfile as sf, pyarrow.parquet as pq |
| os.environ["HF_HUB_OFFLINE"] = "1" |
| print("device", torch.cuda.get_device_name()) |
| MD = os.path.dirname(os.path.abspath(__file__)) |
| MOD, CLS, SRC, CKPT = { |
| "Forensics-WavLM-Base": ("forensics_wavlm_base", "ForensicsWavLMBase", |
| "eliya__forensics_0.3B_base_deepfake_classifier", "checkpoint_epoch_5.safetensors"), |
| "Forensics-WavLM-OCSoftmax": ("forensics_wavlm_ocsoftmax", "ForensicsWavLMOCSoftmax", |
| "eliya__forensics_0.3B_wavlm_oc_softmax_deepfake_classifier", |
| "2nd_oc_softmax_ce_from16_last.safetensors"), |
| }[os.path.basename(MD)] |
| SRC = os.path.join(MD, "..", "_ckpt_dl", SRC) |
| from safetensors.torch import load_file, save_file |
| import transformers |
| def imp(name, path): |
| spec = importlib.util.spec_from_file_location(name, path); m = importlib.util.module_from_spec(spec) |
| sys.modules[name] = m; spec.loader.exec_module(m); return m |
| sd = load_file(f"{SRC}/{CKPT}") |
| tmp = tempfile.mkdtemp() |
| import shutil; shutil.copy(f"{MD}/wavlm_large_config.json", f"{tmp}/config.json") |
| save_file({k[len("wavlm."):]: v.contiguous() for k, v in sd.items() if k.startswith("wavlm.")}, f"{tmp}/model.safetensors") |
| _orig = transformers.WavLMModel.from_pretrained.__func__ |
| transformers.WavLMModel.from_pretrained = classmethod(lambda cls, name, *a, **k: _orig(cls, tmp, *a, **k)) |
| srcmodel = imp("model", f"{SRC}/model.py") |
| srcinf = imp("inference", f"{SRC}/inference.py") |
| ref = srcmodel.DeepfakeDetector().cuda().eval() |
| r = ref.load_state_dict(sd, strict=False) |
| print("source strict=False:", "missing", r.missing_keys, "unexpected", r.unexpected_keys) |
| sys.path.insert(0, MD) |
| F = importlib.import_module(MOD) |
| m = getattr(F, CLS)(); m.load() |
| print("params", sum(p.numel() for p in m.net.parameters())) |
| t = pq.read_table(f"{MD}/../InTheWild/data/test-00000-of-00008.parquet", columns=["audio", "label"]).slice(0, 400).to_pylist() |
| by = {0: [], 1: []} |
| for row in t: |
| if len(by[row["label"]]) < 10: |
| by[row["label"]].append(row["audio"]["bytes"]) |
| a, b, lab = [], [], [] |
| for y in (0, 1): |
| for i, byts in enumerate(by[y]): |
| f = f"{tmp}/clip_{y}_{i}.wav"; open(f, "wb").write(byts) |
| with torch.no_grad(): |
| a.append(ref(srcinf.load_audio(f).unsqueeze(0).cuda()).float().item()) |
| w, sr = sf.read(f, dtype="float32"); assert sr == 16000, sr |
| if w.ndim > 1: w = w.mean(1) |
| b.append(m.score_batch([w], [16000])[0]); lab.append(y) |
| a, b = np.array(a), np.array(b) |
| print("src ", np.round(a, 4)); print("ours", np.round(b, 4)); print("labels", lab) |
| print("n", len(a), "max abs diff", np.abs(a - b).max()) |
| shutil.rmtree(tmp) |
|
|
