import torch
from transformers import AutoModel, AutoTokenizer, WhisperFeatureExtractor

stage = "stage_b"

model = AutoModel.from_pretrained(
    "RaghaRao314159/transcription-models",
    subfolder=stage,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained(
    "RaghaRao314159/transcription-models",
    subfolder=stage,
    trust_remote_code=True,
)
feature_extractor = WhisperFeatureExtractor.from_pretrained(
    "RaghaRao314159/transcription-models",
    subfolder=stage,
)

Compare Both Stages

python pull_model_and_infer.py --model-source "RaghaRao314159/transcription-models" --stage both --audio-path test.mp3