File size: 3,890 Bytes
8c0bdcb 37f59d3 8c0bdcb | 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 111 112 113 114 115 116 117 | """Static types shared by the asr_* modules: audio aliases, loader kwargs, protocols.
Typing scaffolding with no behaviour of its own: the `TypedDict`s describe
the keyword arguments ASRModel forwards to `from_pretrained` / `cached_file`,
and the protocol classes describe surfaces with several implementations (or
test fakes) and no common library base that declares them.
"""
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Protocol, TypedDict, TypeVar
import numpy.typing as npt
import torch
from transformers import BatchFeature, PreTrainedModel
from transformers.generation.utils import GenerateOutput, GenerationMixin
# One waveform (array, tensor or list of samples), or a batch of them.
Waveform = npt.ArrayLike | torch.Tensor
AudioInput = Waveform | Sequence[Waveform]
# A `state_dict(destination=...)` mapping, returned as the caller's own type.
StateDictT = TypeVar("StateDictT", bound=dict[str, Any])
class PreparedChunk(TypedDict):
"""One chunk's encoder inputs, as `ASRPipeline.prepare_chunk` builds them (batch of 1)."""
input_features: torch.Tensor
attention_mask: torch.Tensor
if TYPE_CHECKING:
class GenerativeDecoder(PreTrainedModel, GenerationMixin):
"""Static view of the decoder: a PreTrainedModel that can generate.
transformers annotates GenerationMixin's `self` with a protocol
(`GenerativePreTrainedModel`) that no PreTrainedModel satisfies -- it
asks for a mutable `device` where the model has a read-only property --
so neither checker can bind `generate` on a real model. This restates
the two generation entry points ASRModel calls without that bound.
Type-checking only; the runtime object is whatever the loader returned.
"""
def generate(self, *args: Any, **kwargs: Any) -> GenerateOutput | torch.LongTensor:
"""`GenerationMixin.generate`."""
...
def prepare_inputs_for_generation(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
"""`GenerationMixin.prepare_inputs_for_generation`."""
...
class LoadKwargs(TypedDict):
"""Loader arguments shared by the encoder and decoder `from_pretrained` calls."""
attn_implementation: str | None
low_cpu_mem_usage: bool
dtype: torch.dtype
class DecoderLoadKwargs(LoadKwargs):
"""Decoder loader arguments: the shared ones plus remote-code opt-in."""
trust_remote_code: bool
class HubFileKwargs(TypedDict, total=False):
"""Hub location arguments forwarded to `cached_file` and PEFT's loader."""
subfolder: str
revision: str
class AudioFeatureExtractor(Protocol):
"""A concrete audio feature extractor's `__call__` (e.g. Whisper's).
`SequenceFeatureExtractor` declares no `__call__`, though every concrete
extractor has one; cast an extractor to this to call it.
"""
def __call__(
self,
raw_speech: AudioInput,
/,
*,
sampling_rate: int,
return_attention_mask: bool,
return_tensors: str,
padding: bool | str = ...,
**kwargs: Any,
) -> BatchFeature:
"""Featurize raw audio into `input_features` (plus `attention_mask`)."""
...
class OutputLengthProjector(Protocol):
"""The projector surface the token-count check needs."""
def get_output_length(self, input_length: torch.Tensor) -> torch.Tensor:
"""Projector output length for encoder output length `input_length`."""
...
class LoadStateDictResult(Protocol):
"""What `load_state_dict` reports back (torch's `_IncompatibleKeys`)."""
@property
def missing_keys(self) -> Sequence[str]:
"""Model keys the state dict did not provide."""
...
@property
def unexpected_keys(self) -> Sequence[str]:
"""State-dict keys the model has no slot for."""
...
|