audio_emotion_detector / download_model.py
kedar-bhumkar's picture
Upload 12 files
ae55e39 verified
import os
import torch
from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2FeatureExtractor
# Define the model name
MODEL_NAME = "superb/wav2vec2-base-superb-er"
OUTPUT_DIR = "." # Save directly to current directory
print(f"Downloading model: {MODEL_NAME}")
print("This may take a few minutes depending on your internet connection...")
try:
# Download and save the feature extractor
print("Downloading feature extractor...")
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_NAME)
feature_extractor.save_pretrained(OUTPUT_DIR)
print(f"Feature extractor saved to current directory")
# Download and save the model
print("Downloading model...")
model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_NAME)
model.save_pretrained(OUTPUT_DIR)
print(f"Model saved to current directory")
print("\nModel and feature extractor downloaded successfully!")
print("You can now use them in your application by loading from the current directory.")
except Exception as e:
print(f"Error downloading model: {e}")
print("\nTrying alternative approach...")
# If direct download fails, try using AutoClasses first and then save with specific classes
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
# Download with Auto classes
feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
model = AutoModelForAudioClassification.from_pretrained(MODEL_NAME)
# Save with specific classes
feature_extractor.save_pretrained(OUTPUT_DIR)
model.save_pretrained(OUTPUT_DIR)
print("\nModel and feature extractor downloaded successfully using alternative approach!")
print("You can now use them in your application by loading from the current directory.")