| """ | |
| A script that is run when the server starts. | |
| """ | |
| import os | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| DOWNLOADED_MODELS_DIR = os.path.join(os.path.dirname(__file__), 'models') | |
| def download_model(model_name: str): | |
| """ | |
| Downloads a model from hugging face hub to the disk but not to the RAM. | |
| :param model_name: The name of the model to download. | |
| """ | |
| print(f"Downloading model: {model_name}") | |
| model = AutoModelForCausalLM.from_pretrained(model_name) | |
| del model | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| del tokenizer | |
| print(f"Downloaded model: {model_name}") | |
| def download_useful_models(): | |
| """ | |
| Downloads the models that are useful for this project. | |
| So that the user doesn't have to wait for the models to download when they first use the app. | |
| """ | |
| print("Downloading useful models...") | |
| os.makedirs(DOWNLOADED_MODELS_DIR, exist_ok=True) | |
| useful_models = ( | |
| "gpt2", | |
| "EleutherAI/gpt-j-6B", | |
| "sberbank-ai/mGPT", | |
| "facebook/opt-125m", | |
| ) | |
| for model_name in useful_models: | |
| download_model(model_name) | |
| def main(): | |
| download_useful_models() | |
| if __name__ == "__main__": | |
| main() | |