yonikremer
commited on
Commit
•
0461bfe
1
Parent(s):
b598d9f
deleting the downloaded models from the memory
Browse files- on_server_start.py +12 -3
on_server_start.py
CHANGED
@@ -1,17 +1,24 @@
|
|
1 |
"""
|
2 |
A script that is run when the server starts.
|
3 |
"""
|
|
|
4 |
|
5 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
6 |
|
|
|
|
|
7 |
|
8 |
def download_model(model_name: str):
|
9 |
"""
|
10 |
-
Downloads
|
11 |
:param model_name: The name of the model to download.
|
12 |
"""
|
13 |
-
|
14 |
-
|
|
|
|
|
|
|
|
|
15 |
|
16 |
|
17 |
def download_useful_models():
|
@@ -19,6 +26,8 @@ def download_useful_models():
|
|
19 |
Downloads the models that are useful for this project.
|
20 |
So that the user doesn't have to wait for the models to download when they first use the app.
|
21 |
"""
|
|
|
|
|
22 |
useful_models = (
|
23 |
"gpt2",
|
24 |
"EleutherAI/gpt-j-6B",
|
|
|
1 |
"""
|
2 |
A script that is run when the server starts.
|
3 |
"""
|
4 |
+
import os
|
5 |
|
6 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
7 |
|
8 |
+
DOWNLOADED_MODELS_DIR = os.path.join(os.path.dirname(__file__), 'models')
|
9 |
+
|
10 |
|
11 |
def download_model(model_name: str):
|
12 |
"""
|
13 |
+
Downloads a model from hugging face hub to the disk but not to the RAM.
|
14 |
:param model_name: The name of the model to download.
|
15 |
"""
|
16 |
+
print(f"Downloading model: {model_name}")
|
17 |
+
model = AutoModelForCausalLM.from_pretrained(model_name)
|
18 |
+
del model
|
19 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
20 |
+
del tokenizer
|
21 |
+
print(f"Downloaded model: {model_name}")
|
22 |
|
23 |
|
24 |
def download_useful_models():
|
|
|
26 |
Downloads the models that are useful for this project.
|
27 |
So that the user doesn't have to wait for the models to download when they first use the app.
|
28 |
"""
|
29 |
+
print("Downloading useful models...")
|
30 |
+
os.makedirs(DOWNLOADED_MODELS_DIR, exist_ok=True)
|
31 |
useful_models = (
|
32 |
"gpt2",
|
33 |
"EleutherAI/gpt-j-6B",
|