fix global vars bug

Author: feifeibear

globals.py

@@ -0,0 +1,19 @@
+import torch
+
+class Singleton(type):
+    _instances = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            cls._instances[cls] = super().__call__(*args, **kwargs)
+        return cls._instances[cls]
+
+class Decoder(metaclass=Singleton):
+    def __init__(self):
+        self.tokenizer = None
+
+    def set_tokenizer(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    def decode(self, t: torch.Tensor) -> str:
+        return self.tokenizer.decode(t[0], skip_special_tokens=True)

main.py

@@ -6,16 +6,7 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 from sampling import autoregressive_sampling, speculative_sampling, speculative_sampling_v2
-
-class Decoder:
-    def __init__(self, tokenizer) -> None:
-        self.tokenizer = tokenizer
-    
-    def decode(self, t : torch.Tensor) -> str:
-        # assert t.dim == 2, "t must be 2d tensor"
-        return self.tokenizer.decode(t[0], skip_special_tokens=True)
-
-DECODER : Decoder = None    
+from globals import Decoder
 
 # my local models
 MODELZOO = {
@@ -34,8 +25,8 @@ def parse_arguments():
     parser = argparse.ArgumentParser(description='args for main.py')
 
     parser.add_argument('--input', type=str, default="Suggest at least five related search terms to \"Mạng neural nhân tạo\".")
-    parser.add_argument('--approx_model_name', type=str, default=MODELZOO["llama1b"])
-    parser.add_argument('--target_model_name', type=str, default=MODELZOO["llama7b"])
+    parser.add_argument('--approx_model_name', type=str, default=MODELZOO["bloom-560m"])
+    parser.add_argument('--target_model_name', type=str, default=MODELZOO["bloom7b"])
     parser.add_argument('--verbose', '-v', action='store_true', default=False, help='enable verbose mode')
     parser.add_argument('--seed', '-s', type=int, default=None, help='set a random seed')
     args = parser.parse_args()
@@ -48,9 +39,8 @@ def generate(input_text, approx_model_name, target_model_name, num_tokens=40, ra
     torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     tokenizer = AutoTokenizer.from_pretrained(approx_model_name, trust_remote_code=True)
-    
-    global DECODER
-    DECODER = Decoder(tokenizer)
+  
+    Decoder().set_tokenizer(tokenizer)
 
     print(f"begin loading models: \n {approx_model_name} \n {target_model_name}")
     small_model = AutoModelForCausalLM.from_pretrained(approx_model_name, trust_remote_code=True).to(torch_device)

sampling/speculative_sampling.py

@@ -4,6 +4,7 @@
 
 from sampling.kvcache_model import KVCacheModel
 from sampling.utils import norm_logits, sample, max_fn
+from globals import Decoder
 
 @torch.no_grad()
 def speculative_sampling(prefix : torch.Tensor, approx_model : torch.nn.Module, target_model : torch.nn.Module, 
@@ -37,9 +38,6 @@ def speculative_sampling(prefix : torch.Tensor, approx_model : torch.nn.Module,
 
     device = target_model.device
 
-    if verbose:
-        global DECODER
-    
     approx_model_cache = KVCacheModel(approx_model, temperature, top_k, top_p, random_seed)
     target_model_cache = KVCacheModel(target_model, temperature, top_k, top_p, random_seed)
 
@@ -62,7 +60,7 @@ def speculative_sampling(prefix : torch.Tensor, approx_model : torch.nn.Module,
                 break
 
             if verbose:
-                print(f"approx guess accepted {j[0]}: \033[31m{DECODER.decode(torch.tensor([j]))}\033[0m")
+                print(f"approx guess accepted {j[0]}: \033[31m{Decoder().decode(torch.tensor([j]))}\033[0m")
 
         # print(f"n : {n}, i : {i}, prefix_len + gamma - 1: {prefix_len + gamma - 1}")
         assert n >= prefix_len - 1, f"n {n}, prefix_len {prefix_len}"
@@ -76,15 +74,15 @@ def speculative_sampling(prefix : torch.Tensor, approx_model : torch.nn.Module,
             # reject someone, sample from the pos n
             t = sample(max_fn(target_model_cache._prob_history[:, n, :] - approx_model_cache._prob_history[:, n, :]), random_seed=random_seed)
             if verbose:
-                print(f"target resamples at position {n}: \033[34m{DECODER.decode(t)}\033[0m")
+                print(f"target resamples at position {n}: \033[34m{Decoder().decode(t)}\033[0m")
 
             target_model_cache.rollback(n+1)
         else:
             # all approx model decoding accepted
             assert n == target_model_cache._prob_history.shape[1] - 1
             t = sample(target_model_cache._prob_history[:, -1, :], random_seed=random_seed)
             if verbose:
-                print(f"target samples {n}: \033[35m{DECODER.decode(t)}\033[0m")
+                print(f"target samples {n}: \033[35m{Decoder().decode(t)}\033[0m")
             target_model_cache.rollback(n+2)