Spaces:
Sleeping
Sleeping
File size: 3,786 Bytes
88a679b | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | import os
import sys
import torch
import json
import sounddevice as sd
import numpy as np
from transformers import Wav2Vec2Processor
from src.models.phoneme_embedder import Wav2Vec2PhonemeEmbedder
from src.utils.audio_utils import AudioPreprocessor
def clear_screen():
os.system('cls' if os.name == 'nt' else 'clear')
class LiveTester:
def __init__(self, model_dir):
print(f"Initializing model from {model_dir}...")
self.processor = Wav2Vec2Processor.from_pretrained(model_dir)
self.model = Wav2Vec2PhonemeEmbedder.from_pretrained(model_dir)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
self.model.eval()
vocab_path = os.path.join(model_dir, "vocab.json")
with open(vocab_path, "r", encoding="utf8") as f:
self.vocab = json.load(f)
self.id2phoneme = {v: k for k, v in self.vocab.items()}
self.pad_id = self.processor.tokenizer.pad_token_id
self.audio_prep = AudioPreprocessor(sr=16000)
self.sr = 16000
def record_and_transcribe(self, duration=3.0):
print(f"\n🎤 Recording for {duration} seconds... Speak now!")
recording = sd.rec(int(duration * self.sr), samplerate=self.sr, channels=1)
sd.wait()
print("✅ Recording complete. Processing...")
audio = recording.squeeze()
# Clean audio with VAD/FFT
audio = self.audio_prep.preprocess(audio)
inputs = self.processor(audio, sampling_rate=self.sr, return_tensors="pt", padding=True)
input_values = inputs.input_values.to(self.device)
with torch.no_grad():
logits = self.model(input_values).logits
pred_ids = torch.argmax(logits, dim=-1)[0].cpu().numpy()
collapsed = []
prev = None
for pid in pred_ids:
if pid != prev and pid != self.pad_id:
collapsed.append(self.id2phoneme.get(pid, "<unk>"))
prev = pid
return collapsed
def main():
model_dir = "trained_models/20k_steps"
if not os.path.exists(model_dir):
print(f"Error: {model_dir} not found. Adjust the path in cli_test_menu.py")
return
tester = LiveTester(model_dir)
while True:
clear_screen()
print("="*50)
print("🎙️ PHONEME EMBEDDER LIVE TEST MENU")
print("="*50)
print(f"Model: {model_dir}")
print("-"*50)
print("1. Quick Test (3 seconds)")
print("2. Long Test (5 seconds)")
print("3. Custom Duration")
print("4. Exit")
print("-"*50)
choice = input("Select an option: ")
if choice == '1':
result = tester.record_and_transcribe(3.0)
print("\nPREDICTED PHONEMES:")
print(" ".join(result))
input("\nPress Enter to continue...")
elif choice == '2':
result = tester.record_and_transcribe(5.0)
print("\nPREDICTED PHONEMES:")
print(" ".join(result))
input("\nPress Enter to continue...")
elif choice == '3':
try:
dur = float(input("Enter duration in seconds: "))
result = tester.record_and_transcribe(dur)
print("\nPREDICTED PHONEMES:")
print(" ".join(result))
except ValueError:
print("Invalid duration.")
input("\nPress Enter to continue...")
elif choice == '4':
print("Goodbye!")
break
else:
print("Invalid choice.")
input("\nPress Enter to continue...")
if __name__ == "__main__":
main()
|