Spaces:
Sleeping
Sleeping
File size: 1,468 Bytes
82e11f0 | 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 | from typing import cast
import torch
from pydantic import BaseModel, ConfigDict
from transformers import (
AutoFeatureExtractor,
AutoTokenizer,
PreTrainedTokenizerBase,
SequenceFeatureExtractor,
)
from src.modeling.config import ModelConfig
from src.modeling.lambda_audio import LambdaAudioModel
from src.modeling.lfm2_decoder import AudioLfm2ForCausalLM
from src.modeling.zipformer import ZipformerModel
class ModelComponents(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
model: LambdaAudioModel
feature_extractor: SequenceFeatureExtractor
tokenizer: PreTrainedTokenizerBase
def load_model_components(
config: ModelConfig,
token: str,
) -> ModelComponents:
feature_extractor = cast(
SequenceFeatureExtractor,
AutoFeatureExtractor.from_pretrained(config.encoder_model_id, token=token),
)
tokenizer = cast(
PreTrainedTokenizerBase,
AutoTokenizer.from_pretrained(config.decoder_model_id, token=token),
)
encoder = ZipformerModel.from_pretrained(config.encoder_model_id, token=token)
decoder = AudioLfm2ForCausalLM.from_pretrained(
config.decoder_model_id,
attn_implementation="sdpa",
dtype=torch.float32,
token=token,
)
model = LambdaAudioModel(encoder, decoder)
return ModelComponents(
model=model,
feature_extractor=feature_extractor,
tokenizer=tokenizer,
)
|