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.")