Delete .ipynb_checkpoints/infer-checkpoint.py with huggingface_hub

.ipynb_checkpoints/infer-checkpoint.py

@@ -1,168 +0,0 @@
-#!/usr/bin/env python3
-"""
-Standalone inference script for the best sdclap CLAP model (qwen3 text encoder + Whisper-large-v3
-last-layer audio encoder, joint_dim=256, mlp3_residual projection heads). Computes a single
-cosine-similarity score between one audio clip and one text description of how it was spoken.
-
-This script is self-contained on purpose — it does NOT import anything from the main sdclap
-repository (`scripts/neural_network/...`); the minimal model/embedding logic it needs is
-reimplemented inline below, so it only depends on the packages in `requirements.txt` plus a
-checkpoint file from `checkpoints/`.
-
-The checkpoint (~11MB) contains ONLY the trained projection heads + gender embedding table +
-logit_scale/bias — NOT the frozen pretrained text/audio encoders, which are downloaded from
-HuggingFace on first run (Qwen/Qwen3-Embedding-4B ~8GB, openai/whisper-large-v3 ~3GB).
-
-Usage (venv with requirements.txt installed):
-    python infer.py --audio clip.wav --text "тихо прошептала, едва сдерживая слёзы"
-    python infer.py --audio clip.wav --text "..." --gender male --checkpoint checkpoints/qwen3_whisper_last_seed2_best.pt
-"""
-from __future__ import annotations
-
-import argparse
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-GENDER_TO_ID = {"male": 0, "female": 1, "unknown": 2}
-TEXT_ENCODER_NAME = "Qwen/Qwen3-Embedding-4B"
-AUDIO_ENCODER_NAME = "openai/whisper-large-v3"
-WHISPER_LAYER_IDX = 32  # last encoder layer of whisper-large-v3 ("whisper_large_v3_l32")
-WHISPER_ENCODER_FPS = 50.0
-WHISPER_MAX_FRAMES = 1500
-
-
-# --- minimal ClapModel (inference-only path: no VIB, no FiLM, no n_vectors>1, no attention-pool —
-# none of these are used by the qwen3_whisper_last checkpoints published here) --------------------
-
-class _ResidualBlock(nn.Module):
-    def __init__(self, dim: int, dropout: float) -> None:
-        super().__init__()
-        self.block = nn.Sequential(
-            nn.Linear(dim, dim), nn.LayerNorm(dim), nn.GELU(), nn.Dropout(dropout),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return x + self.block(x)
-
-
-class ProjectionHead(nn.Module):
-    """mlp3_residual architecture only — matches head_arch used by qwen3_whisper_last."""
-
-    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, dropout: float) -> None:
-        super().__init__()
-        self.backbone = nn.Sequential(
-            nn.Linear(in_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU(), nn.Dropout(dropout),
-            _ResidualBlock(hidden_dim, dropout),
-        )
-        self.out_head = nn.Linear(hidden_dim, out_dim)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.out_head(self.backbone(x))
-
-
-class ClapModel(nn.Module):
-    def __init__(self, text_dim: int, audio_dim: int, joint_dim: int, hidden_dim: int,
-                 gender_dim: int, dropout: float) -> None:
-        super().__init__()
-        self.gender_embedding = nn.Embedding(3, gender_dim)
-        self.text_head = ProjectionHead(text_dim + gender_dim, hidden_dim, joint_dim, dropout)
-        self.audio_head = ProjectionHead(audio_dim, hidden_dim, joint_dim, dropout)
-
-    def encode_text(self, text_emb: torch.Tensor, gender_id: torch.Tensor) -> torch.Tensor:
-        gender_vec = self.gender_embedding(gender_id)
-        x = torch.cat([text_emb, gender_vec], dim=-1)
-        return F.normalize(self.text_head(x), dim=-1)
-
-    def encode_audio(self, audio_emb: torch.Tensor) -> torch.Tensor:
-        return F.normalize(self.audio_head(audio_emb), dim=-1)
-
-
-def load_clap_model(checkpoint_path: str, device: str) -> ClapModel:
-    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
-    sd = ckpt["model"]
-    args = ckpt["args"]
-    text_dim = sd["text_head.backbone.0.weight"].shape[1] - args["gender_dim"]
-    audio_dim = sd["audio_head.backbone.0.weight"].shape[1]
-    model = ClapModel(
-        text_dim=text_dim, audio_dim=audio_dim, joint_dim=args["joint_dim"],
-        hidden_dim=args["hidden_dim"], gender_dim=args["gender_dim"], dropout=args["dropout"],
-    )
-    missing, unexpected = model.load_state_dict(sd, strict=False)
-    unexpected = [k for k in unexpected if k not in ("logit_scale", "logit_bias")]
-    if missing or unexpected:
-        raise RuntimeError(f"checkpoint key mismatch: missing={missing}, unexpected={unexpected}")
-    model.to(device).eval()
-    return model
-
-
-# --- raw text -> qwen3 embedding (2560-dim, L2-normalized) ----------------------------------------
-
-def embed_text(text: str, device: str) -> torch.Tensor:
-    from sentence_transformers import SentenceTransformer
-
-    kwargs = {
-        "model_kwargs": {"dtype": torch.bfloat16 if device == "cuda" else torch.float32},
-        "processor_kwargs": {"padding_side": "left"},
-    }
-    encoder = SentenceTransformer(TEXT_ENCODER_NAME, device=device, **kwargs)
-    emb = encoder.encode([text], normalize_embeddings=True, show_progress_bar=False)
-    return torch.from_numpy(np.asarray(emb, dtype=np.float32)).to(device)
-
-
-# --- raw audio -> whisper-large-v3 last-layer embedding (1280-dim, masked mean-pool) --------------
-
-def embed_audio(audio_path: str, device: str) -> torch.Tensor:
-    import librosa
-    from transformers import WhisperFeatureExtractor, WhisperModel
-
-    extractor = WhisperFeatureExtractor.from_pretrained(AUDIO_ENCODER_NAME)
-    model = WhisperModel.from_pretrained(AUDIO_ENCODER_NAME, dtype=torch.float16).to(device).eval()
-    encoder = model.encoder
-    target_sr = extractor.sampling_rate
-
-    wave, sr = librosa.load(audio_path, sr=None, mono=True)
-    if sr != target_sr:
-        wave = librosa.resample(wave.astype(np.float32), orig_sr=sr, target_sr=target_sr)
-    duration_s = len(wave) / target_sr
-    n_real = min(WHISPER_MAX_FRAMES, max(1, round(duration_s * WHISPER_ENCODER_FPS)))
-
-    inputs = extractor(wave, sampling_rate=target_sr, return_tensors="pt")
-    input_features = inputs["input_features"].to(device=device, dtype=torch.float16)
-    with torch.no_grad():
-        out = encoder(input_features, output_hidden_states=True)
-    hs = out.hidden_states[WHISPER_LAYER_IDX][:, :n_real, :]
-    audio_vec = hs.mean(dim=1)[0].float()
-    return audio_vec.unsqueeze(0).to(device)
-
-
-def main() -> None:
-    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    ap.add_argument("--audio", required=True, help="path to an audio file (wav/mp3/flac/...)")
-    ap.add_argument("--text", required=True, help="Russian description of how the speech sounds")
-    ap.add_argument("--gender", default="unknown", choices=list(GENDER_TO_ID),
-                     help="reader gender to condition on — 'unknown' matches the paper's zero-shot "
-                          "evaluation setup on external datasets (default)")
-    ap.add_argument("--checkpoint", default="checkpoints/qwen3_whisper_last_seed0_best.pt",
-                     help="any of the 5 published seeds are equivalent for this purpose")
-    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
-    args = ap.parse_args()
-
-    model = load_clap_model(args.checkpoint, args.device)
-
-    text_emb = embed_text(args.text, args.device)
-    audio_emb = embed_audio(args.audio, args.device)
-
-    gender_id = torch.tensor([GENDER_TO_ID[args.gender]], dtype=torch.long, device=args.device)
-    with torch.no_grad():
-        text_z = model.encode_text(text_emb, gender_id)
-        audio_z = model.encode_audio(audio_emb)
-        similarity = (text_z * audio_z).sum(dim=-1).item()
-
-    print(f"cosine similarity: {similarity:.4f}")
-
-
-if __name__ == "__main__":
-    main()