import os
from sentence_transformers import SentenceTransformer

# Configuration
MODEL_NAME = 'BAAI/bge-base-en-v1.5'  # The 768-dimension model
SAVE_PATH = './models/bge-base-en-v1.5'

def download_model():
    """Download and save the embedding model locally."""
    print(f"Downloading model: {MODEL_NAME}...")
    
    # Download and load the model
    model = SentenceTransformer(MODEL_NAME)
    
    # Save it to the specific folder
    os.makedirs(SAVE_PATH, exist_ok=True)
    print(f"Saving model to: {SAVE_PATH}...")
    model.save(SAVE_PATH)
    
    print("✅ Model downloaded and saved successfully.")
    
    # Check model file size
    model_file = os.path.join(SAVE_PATH, 'model.safetensors')
    if os.path.exists(model_file):
        size_mb = os.path.getsize(model_file) / (1024 * 1024)
        print(f"Model file size: {size_mb:.2f} MB")
    
    print(f"Model dimension: {model.get_sentence_embedding_dimension()}")

if __name__ == "__main__":
    download_model()