'''
Downloads models from Hugging Face to models/model-name.

Example:
python download-model.py facebook/opt-1.3b

'''

import requests 
from bs4 import BeautifulSoup 
import multiprocessing
import tqdm
from sys import argv
from pathlib import Path

def get_file(args):
    url = args[0]
    output_folder = args[1]

    r = requests.get(url, stream=True)
    with open(output_folder / Path(url.split('/')[-1]), 'wb') as f:
        total_size = int(r.headers.get('content-length', 0))
        block_size = 1024
        t = tqdm.tqdm(total=total_size, unit='iB', unit_scale=True)
        for data in r.iter_content(block_size):
            t.update(len(data))
            f.write(data)
        t.close()

if __name__ == '__main__':
    model = argv[1]
    if model[-1] == '/':
        model = model[:-1]
    url = f'https://huggingface.co/{model}/tree/main'
    output_folder = Path("models") / model.split('/')[-1]
    if not output_folder.exists():
        output_folder.mkdir()

    # Finding the relevant files to download
    page = requests.get(url) 
    soup = BeautifulSoup(page.content, 'html.parser') 
    links = soup.find_all('a')
    downloads = []
    for link in links:
        href = link.get('href')[1:]
        if href.startswith(f'{model}/resolve/main'):
            if href.endswith(('.json', '.txt')) or (href.endswith('.bin') and 'pytorch_model' in href):
                downloads.append(f'https://huggingface.co/{href}')

    # Downloading the files
    print(f"Downloading the model to {output_folder}...")
    pool = multiprocessing.Pool(processes=4)
    results = pool.map(get_file, [[downloads[i], output_folder] for i in range(len(downloads))])
    pool.close()
    pool.join()