File size: 1,572 Bytes
a4534ab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | #!/usr/bin/env python3
"""Generate a PROPS speaker-embedding GMM from a natural-language description."""
from __future__ import annotations
import argparse
from pathlib import Path
import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from model import ComposedGMMMDN
def main():
parser = argparse.ArgumentParser()
parser.add_argument("description", help="Natural-language speaker-profile description")
parser.add_argument("--model", default=str(Path(__file__).resolve().parent))
parser.add_argument("--output", type=Path, default=Path("generated_gmm.npz"))
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
args = parser.parse_args()
encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device=args.device)
embedding = encoder.encode([args.description], convert_to_tensor=True).to(args.device)
model = ComposedGMMMDN.from_pretrained(args.model, device=args.device)
with torch.inference_mode():
pi_logits, mu, sigma = model(embedding)
pi = torch.softmax(pi_logits, dim=-1)
args.output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(
args.output,
pi_logits=pi_logits[0].cpu().numpy().astype(np.float32),
pi=pi[0].cpu().numpy().astype(np.float32),
mu=mu[0].cpu().numpy().astype(np.float32),
sigma=sigma[0].cpu().numpy().astype(np.float32),
description=np.asarray(args.description),
)
print(args.output)
if __name__ == "__main__":
main()
|