ml/precache_models.py

"""
ml/precache_models.py - Pre-cache Hugging Face transformer models into Docker container
======================================================================================
Pre-downloads base weights during docker build to ensure sub-second startup on Render.
"""
import sys

def main():
    print("[Pre-Cache] Pre-caching base transformer models into container image...")
    try:
        from transformers import (
            Wav2Vec2ForSequenceClassification,
            Wav2Vec2FeatureExtractor,
            Wav2Vec2ForCTC,
            Wav2Vec2Processor,
        )
        Wav2Vec2ForSequenceClassification.from_pretrained("facebook/wav2vec2-base")
        Wav2Vec2FeatureExtractor.from_pretrained("facebook/wav2vec2-base")
        Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
        Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
        print("[Pre-Cache] Base models successfully cached.")
    except Exception as e:
        print(f"[Pre-Cache] Warning during pre-caching: {e}", file=sys.stderr)

if __name__ == "__main__":
    main()