| """VeriSpoof network: audio encoder -> audio tokenizer -> language model, plus a per-token head.""" |
| import torch |
| import torch.nn as nn |
| from safetensors.torch import load_file |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| from verispoof.audio_tokenizer import AudioTokenizer |
| from verispoof.encoder import AudioEncoder |
|
|
| AUDIO_TOKEN = "<unused_token_0>" |
| SYSTEM = ("You are VeriSpoof, an assistant that analyses speech recordings for deepfake detection and for " |
| "understanding how synthetic speech was generated.") |
| AUDIO = "<audio>" |
| END = "<|im_end|>" |
|
|
|
|
| class Network(nn.Module): |
| def __init__(self, path, cfg, device="cuda", load_llm=True): |
| """load_llm=False: only the audio front end and the input embedding table (for an external LLM engine).""" |
| super().__init__() |
| self.time_marks = cfg.get("time_marks") |
| self.system = cfg.get("system_prompt", SYSTEM) |
| self.window = cfg.get("window", 128000) |
| dtype = getattr(torch, cfg.get("llm_dtype", "bfloat16")) |
| self.amp = dtype != torch.float32 |
| self.text_tok = AutoTokenizer.from_pretrained(f"{path}/llm") |
| if load_llm: |
| self.llm = AutoModelForCausalLM.from_pretrained(f"{path}/llm", dtype=dtype, attn_implementation="sdpa") |
| self.llm.requires_grad_(False) |
| self.embed_tokens = self.llm.get_input_embeddings() |
| h = self.llm.config.hidden_size |
| else: |
| from transformers import AutoConfig |
| self.llm = None |
| from safetensors import safe_open |
| with safe_open(f"{path}/llm/model.safetensors", "pt") as f: |
| w = f.get_tensor("model.embed_tokens.weight") |
| self.embed_tokens = nn.Embedding.from_pretrained(w.to(dtype), freeze=True) |
| h = AutoConfig.from_pretrained(f"{path}/llm").hidden_size |
| self.audio_str = AUDIO_TOKEN |
| self.audio_id = self.text_tok.convert_tokens_to_ids(AUDIO_TOKEN) |
| self.audio_tok = AudioTokenizer(h, **cfg["tokenizer"]) |
| self.audio_tok.load_state_dict(load_file(f"{path}/audio_tokenizer.safetensors")) |
| self.tok_head = nn.Sequential(nn.LayerNorm(h), nn.Linear(h, 1)) |
| self.tok_head.load_state_dict(load_file(f"{path}/token_head.safetensors")) |
| self.encoder = AudioEncoder(f"{path}/audio_encoder.safetensors", device=device) |
| self.pad_id = self.text_tok.pad_token_id if self.text_tok.pad_token_id is not None else self.text_tok.eos_token_id |
| self.end_id = self.text_tok.convert_tokens_to_ids(END) |
|
|
| def causal_lm(self): |
| return self.llm |
|
|
| def audio_layout(self, frame_s=0.08): |
| a = self.audio_tok |
| if not self.time_marks: |
| return [a.n_tokens] |
| step = frame_s * a.frame_stack |
| k = max(1, round(self.time_marks / step)) |
| out = [a.n_global] |
| for i in range(0, a.n_local, k): |
| out += [f"[{i * step:.1f}s]", min(k, a.n_local - i)] |
| return out |
|
|
| def prompt_ids(self, prompt): |
| block = "".join(self.audio_str * x if isinstance(x, int) else x for x in self.audio_layout()) |
| text = self.text_tok.apply_chat_template( |
| [{"role": "system", "content": self.system}, {"role": "user", "content": prompt.replace(AUDIO, block)}], |
| tokenize=False, add_generation_prompt=True, enable_thinking=False) |
| return self.text_tok(text, add_special_tokens=False)["input_ids"] |
|
|
| def text_ids(self, text): |
| return self.text_tok(text, add_special_tokens=False)["input_ids"] |
|
|
| def embed(self, input_ids, wav): |
| emb = self.embed_tokens(input_ids) |
| mask = input_ids == self.audio_id |
| hs = self.encoder(wav) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| a = self.audio_tok(hs) |
| assert mask.sum() == a.shape[0] * a.shape[1], "audio placeholders do not match audio inputs" |
| return emb.masked_scatter(mask.unsqueeze(-1), a.to(emb.dtype).reshape(-1, emb.size(-1))) |
|
|
| def token_logits(self, h, input_ids): |
| n = self.audio_tok.n_tokens |
| a = h[input_ids == self.audio_id].view(-1, n, h.size(-1))[:, self.audio_tok.n_global:] |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| return self.tok_head(a).squeeze(-1).float() |
|
|
| def seconds_per_token(self, frame_s=0.08): |
| return frame_s * self.audio_tok.frame_stack |
|
|
