Spaces:
Sleeping
Sleeping
kz209
commited on
Commit
•
9779cd8
1
Parent(s):
e2bb507
add log
Browse files- pages/summarization_playground.py +3 -0
- utils/model.py +3 -0
pages/summarization_playground.py
CHANGED
@@ -8,6 +8,8 @@ from utils.data import dataset
|
|
8 |
import gc
|
9 |
import torch
|
10 |
|
|
|
|
|
11 |
load_dotenv()
|
12 |
|
13 |
__model_on_gpu__ = ''
|
@@ -37,6 +39,7 @@ def generate_answer(sources, model_name, prompt):
|
|
37 |
|
38 |
if __model_on_gpu__ != model_name:
|
39 |
if __model_on_gpu__:
|
|
|
40 |
del model[__model_on_gpu__]
|
41 |
gc.collect()
|
42 |
torch.cuda.empty_cache()
|
|
|
8 |
import gc
|
9 |
import torch
|
10 |
|
11 |
+
import logging
|
12 |
+
|
13 |
load_dotenv()
|
14 |
|
15 |
__model_on_gpu__ = ''
|
|
|
39 |
|
40 |
if __model_on_gpu__ != model_name:
|
41 |
if __model_on_gpu__:
|
42 |
+
logging.info(f"delete model {__model_on_gpu__}")
|
43 |
del model[__model_on_gpu__]
|
44 |
gc.collect()
|
45 |
torch.cuda.empty_cache()
|
utils/model.py
CHANGED
@@ -5,6 +5,8 @@ import torch
|
|
5 |
from huggingface_hub import login
|
6 |
import os
|
7 |
|
|
|
|
|
8 |
login(token = os.getenv('HF_TOKEN'))
|
9 |
|
10 |
class Model(torch.nn.Module):
|
@@ -28,6 +30,7 @@ class Model(torch.nn.Module):
|
|
28 |
device_map="auto",
|
29 |
)
|
30 |
|
|
|
31 |
self.update()
|
32 |
|
33 |
@classmethod
|
|
|
5 |
from huggingface_hub import login
|
6 |
import os
|
7 |
|
8 |
+
import logging
|
9 |
+
|
10 |
login(token = os.getenv('HF_TOKEN'))
|
11 |
|
12 |
class Model(torch.nn.Module):
|
|
|
30 |
device_map="auto",
|
31 |
)
|
32 |
|
33 |
+
logging.info(f'Load model {self.name}')
|
34 |
self.update()
|
35 |
|
36 |
@classmethod
|