From ee5855135eaa13e1fb0743c96ce7aaebf4b624b0 Mon Sep 17 00:00:00 2001 From: Bishwendu Kundu Date: Sun, 23 Oct 2022 17:55:44 +0530 Subject: [PATCH] use gpt-neox model --- download.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/download.py b/download.py index 9f2956d..e76a71f 100644 --- a/download.py +++ b/download.py @@ -1,13 +1,22 @@ # In this file, we define download_model # It runs during container build time to get model weights built into the container -# In this example: A Huggingface BERT model +# In this example: A Huggingface GPT-NeoX model -from transformers import pipeline +from transformers import GPTNeoXForCausalLM, GPTNeoXTokenizerFast +import torch def download_model(): # do a dry run of loading the huggingface model, which will download weights - pipeline('fill-mask', model='bert-base-uncased') + print("downloading model...") + GPTNeoXForCausalLM.from_pretrained( + "EleutherAI/gpt-neox-20b", revision="float16", torch_dtype=torch.float16, low_cpu_mem_usage=True + ) + print("done") + + print("downloading tokenizer...") + GPTNeoXTokenizerFast.from_pretrained("EleutherAI/gpt-neox-20b") + print("done") if __name__ == "__main__": download_model() \ No newline at end of file