TaahKay's picture
Update models/model_loader.py
d329f57 verified
raw
history blame
956 Bytes
import os
import streamlit as st
CACHE_DIR = "/tmp/pretrained_models"
from speechbrain.pretrained.interfaces import foreign_class
from faster_whisper import WhisperModel
# Ensure the folder exists and is writable
os.makedirs(CACHE_DIR, exist_ok=True)
# -------------------------------
# Load Model (Cached)
# -------------------------------
def load_accent_model():
"""Loads custom accent classification model."""
if not os.getenv("HF_TOKEN"):
st.error("Hugging Face token not found.")
st.stop()
try:
return foreign_class(
source="Jzuluaga/accent-id-commonaccent_xlsr-en-english",
pymodule_file="custom_interface.py",
classname="CustomEncoderWav2vec2Classifier",
)
except Exception as e:
st.error(f"❌ Error loading model: {e}")
st.stop()
def load_whisper():
return WhisperModel("tiny", device="cpu", compute_type="int8_float32")