File size: 5,346 Bytes
e71ac9a | 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 | """Waveform preparation and symbolic decoding for SheetSage2."""
import numbers
import torch
import torchaudio
from transformers import ProcessorMixin
from transformers.feature_extraction_utils import BatchFeature
from .tokenization_sheetsage2 import SheetSage2Tokenizer
def _hf_relative_dependencies():
# Keep standalone AutoProcessor loading compatible with Transformers 4.45.
from .durations_sheetsage2 import DURATION_TEMPLATES
from .labels_sheetsage2 import STRUCTURE_LABELS
from .schema_sheetsage2 import get_prompt_multitask_schema
class SheetSage2Processor(ProcessorMixin):
"""Prepare mono waveforms without amplitude normalization.
Pass one mono waveform, a batch tensor, or a list of mono waveforms. Use
``model.transcribe`` for files and audio longer than one model window.
"""
attributes = []
valid_kwargs = ["sampling_rate", "window_seconds", "time_hz", "schema_version", "tokenizer_fingerprint"]
def __init__(self, sampling_rate=24000, window_seconds=300.0, time_hz=100,
schema_version="v1", tokenizer_fingerprint="5ba3325af0344c7f", **kwargs):
super().__init__(**{key: value for key, value in kwargs.items() if key == "chat_template"})
self.sampling_rate = int(sampling_rate)
self.window_seconds = float(window_seconds)
self.time_hz = int(time_hz)
self.schema_version = str(schema_version)
self.tokenizer_fingerprint = str(tokenizer_fingerprint)
self.tokenizer = SheetSage2Tokenizer(
self.window_seconds, self.time_hz, self.schema_version,
expected_fingerprint=self.tokenizer_fingerprint,
)
@property
def model_input_names(self):
return ["input_values", "attention_mask"]
@classmethod
def from_model_config(cls, config):
return cls(config.sampling_rate, config.input_audio_length, config.time_hz,
config.tokenizer_schema_version, config.tokenizer_fingerprint)
def __call__(self, audio, sampling_rate=None, padding=True, return_tensors="pt",
return_attention_mask=True):
source_rate = self.sampling_rate if sampling_rate is None else int(sampling_rate)
if source_rate <= 0:
raise ValueError("sampling_rate must be positive.")
if isinstance(audio, (str, bytes)):
raise ValueError("Pass waveform samples here; use model.transcribe for audio files.")
if isinstance(audio, (list, tuple)) and audio and not isinstance(audio[0], numbers.Number):
waveforms = [torch.as_tensor(value) for value in audio]
else:
value = torch.as_tensor(audio)
waveforms = list(value) if value.ndim == 2 else [value]
if not waveforms:
raise ValueError("Provide at least one waveform.")
prepared = []
maximum = round(self.window_seconds * self.sampling_rate)
for waveform in waveforms:
if waveform.ndim != 1 or not waveform.is_floating_point() or not torch.isfinite(waveform).all():
raise ValueError("Each waveform must be a one-dimensional finite floating-point array.")
waveform = waveform.to(dtype=torch.float32)
if source_rate != self.sampling_rate:
waveform = torchaudio.functional.resample(waveform, source_rate, self.sampling_rate)
if waveform.numel() < 1025:
raise ValueError("Each waveform must contain at least 1025 samples at 24 kHz.")
if waveform.numel() > maximum:
raise ValueError("Audio exceeds one model window; use model.transcribe for whole songs.")
prepared.append(waveform)
if len({str(value.device) for value in prepared}) != 1:
raise ValueError("All waveforms in a batch must be on the same device.")
if padding is False:
length = max(value.numel() for value in prepared)
if any(value.numel() != length for value in prepared):
raise ValueError("Use padding=True for waveforms of different lengths.")
elif padding is True or padding == "max_length":
length = maximum
elif padding == "longest":
length = max(value.numel() for value in prepared)
else:
raise ValueError("padding must be True, False, 'max_length', or 'longest'.")
lengths = torch.tensor([value.numel() for value in prepared], device=prepared[0].device)
values = torch.stack([torch.nn.functional.pad(value, (0, length - value.numel())) for value in prepared])
data = {"input_values": values}
if return_attention_mask:
data["attention_mask"] = (torch.arange(length, device=values.device)[None] < lengths[:, None]).long()
if return_tensors == "np":
data = {key: value.cpu().numpy() for key, value in data.items()}
elif return_tensors not in {None, "pt"}:
raise ValueError("return_tensors must be 'pt', 'np', or None.")
return BatchFeature(data=data)
def decode(self, token_ids, strict=True):
return self.tokenizer.decode_sequence(token_ids, strict=strict)
def batch_decode(self, sequences, strict=True):
return [self.decode(tokens, strict=strict) for tokens in sequences]
SheetSage2Processor.register_for_auto_class()
|