update model_loader

models/extensions/callback.py

@@ -1,9 +1,9 @@
-import gc
+# import gc
 import traceback
 from queue import Queue
-from threading import Thread
+# from threading import Thread
-import threading
+# import threading
-from typing import Optional, List, Dict, Any
+from typing import Optional, List, Dict, Any, TypeVar, Deque
 from collections import deque
 import torch
 import transformers
@@ -12,13 +12,16 @@ from models.extensions.thread_with_exception import ThreadWithException
 import models.shared as shared
 
 
-class LimitedLengthDict(dict):
+K = TypeVar('K')
+V = TypeVar('V')
+
+class LimitedLengthDict(Dict[K, V]):
     def __init__(self, maxlen=None, *args, **kwargs):
         self.maxlen = maxlen
-        self._keys = deque()
+        self._keys: Deque[K] = deque()
         super().__init__(*args, **kwargs)
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: K, value: V):
         if key not in self:
             if self.maxlen is not None and len(self) >= self.maxlen:
                 oldest_key = self._keys.popleft()

models/loader/loader.py

@@ -139,9 +139,9 @@ class LoaderCheckPoint:
 
                     model = dispatch_model(model, device_map=self.device_map)
             else:
-                print(
+                # print(
-                    "Warning: torch.cuda.is_available() returned False.\nThis means that no GPU has been "
+                #     "Warning: torch.cuda.is_available() returned False.\nThis means that no GPU has been "
-                    "detected.\nFalling back to CPU mode.\n")
+                #     "detected.\nFalling back to CPU mode.\n")
                 model = (
                     AutoModel.from_pretrained(
                         checkpoint,