DocClassify / download_model.py
Seth
Update
f6e574f
#!/usr/bin/env python3
"""Download BERT-tiny model upfront to Model folder."""
from pathlib import Path
from transformers import AutoTokenizer, AutoModel
MODEL_NAME = "prajjwal1/bert-tiny"
MODELS_DIR = Path(__file__).resolve().parent / "Model"
MODEL_PATH = MODELS_DIR / "bert-tiny"
def download_model():
"""Download and save the BERT-tiny model."""
print(f"Downloading model: {MODEL_NAME}")
print(f"Target directory: {MODEL_PATH}")
# Create Model directory
MODELS_DIR.mkdir(parents=True, exist_ok=True)
if MODEL_PATH.exists():
print(f"Model already exists at {MODEL_PATH}")
print("Skipping download.")
return
try:
print("Downloading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
print("Downloading model...")
model = AutoModel.from_pretrained(MODEL_NAME)
print(f"Saving model to {MODEL_PATH}...")
tokenizer.save_pretrained(str(MODEL_PATH))
model.save_pretrained(str(MODEL_PATH))
print("✅ Model downloaded and saved successfully!")
print(f"Location: {MODEL_PATH}")
except Exception as e:
print(f"❌ Error downloading model: {e}")
raise
if __name__ == "__main__":
download_model()