Download weights from GECToR pre-trained model

download_weights.py

from pathlib import Path
import requests
from tqdm import tqdm

TEST_FIXTURES_DIR_PATH = "test_fixtures"
MODEL_URL = "https://grammarly-nlp-data-public.s3.amazonaws.com/gector/roberta_1_gectorv2.th"


def download_weights():
    model_path = Path("test_fixtures/roberta_model/weights.th")
    if not model_path.exists():
        response = requests.get(MODEL_URL)
        with model_path.open("wb") as out_fp:
            # Write out data with progress bar
            for data in tqdm(response.iter_content()):
                out_fp.write(data)
    assert model_path.exists()
    return model_path

download_weights()