kz209 commited on
Commit
9779cd8
1 Parent(s): e2bb507
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