fix slow inference

app.py

@@ -15,15 +15,14 @@ import asyncio
 
 tk = None
 model_states = None
-model = None
 lock_using_model = threading.Lock()
 recent_generate_timestamp = time.time()
 
-MODEL_LIVE_TIME = 15 * 60 # 15 minutes
+MODEL_LIVE_TIME = 5#15 * 60 # 15 minutes
 
 
 def load_model():
-    global tk, model_states, model
+    global tk, model_states
 
     tk = transformers.AutoTokenizer.from_pretrained("./tk")
     model_path = "./ts100-re2-h1-4000-model.pt"
@@ -44,18 +43,6 @@ def load_model():
             new_key = k[len(unwanted_prefix_dueto_compile):]
             model_states[new_key] = model_states.pop(k)
     
-    yml_path = "./ts100-re2-h1.yml"
-    with open(yml_path, "r") as f:
-        config = yaml.load(f, Loader=yaml.FullLoader)['model']
-    if "vocab_size" not in config:
-        config['vocab_size'] = tk.vocab_size
-    config["attn_window_size"] = 500
-    # print(config)
-    config = DearthConfig(**config)
-    model = DearthForCausalLM(config)
-
-    model.load_state_dict(model_states)
-    model.eval()
 
 
 def main_free_mem():
@@ -66,12 +53,11 @@ def main_free_mem():
 
 
 def free_mem():
-    global tk, model_states, model, recent_generate_timestamp, lock_using_model
+    global tk, model_states, recent_generate_timestamp, lock_using_model
     lock_using_model.acquire()
-    if time.time() - recent_generate_timestamp >= MODEL_LIVE_TIME and model is not None:
+    if time.time() - recent_generate_timestamp >= MODEL_LIVE_TIME and tk is not None:
         tk = None
         model_states = None
-        model = None
         print(f"free mem, {time.time()}")
     lock_using_model.release()
     try:
@@ -85,13 +71,25 @@ def generate(input, num_more_tokens):
     global tk, model_states, model, recent_generate_timestamp, lock_using_model
     lock_using_model.acquire()
     time_start = time.time()
-    if model is None:
+    if tk is None:
         load_model()
     elif time.time() - recent_generate_timestamp > MODEL_LIVE_TIME:
         tk = None
         model_states = None
-        model = None
         load_model()
+
+    yml_path = "./ts100-re2-h1.yml"
+    with open(yml_path, "r") as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)['model']
+    if "vocab_size" not in config:
+        config['vocab_size'] = tk.vocab_size
+    config["attn_window_size"] = 500
+    # print(config)
+    config = DearthConfig(**config)
+    model = DearthForCausalLM(config)
+
+    model.load_state_dict(model_states)
+    model.eval()
     recent_generate_timestamp = time.time()
     print(f"load model time: {time.time() - time_start}")
 
@@ -158,7 +156,7 @@ if __name__ == "__main__":
         with gr.Row():
             with gr.Column():
                 inp = gr.Textbox(lines=5, label="Input Text", value=example_input[random.randint(0, len(example_input)-1)], elem_id="input_textbox")
-                generate_max_slider = gr.Slider(16, 64, step=1.0, value=32, label="more tokens", info="")
+                generate_max_slider = gr.Slider(8, 64, step=1.0, value=16, label="more tokens", info="")
                 generate_button = gr.Button(value="Generate")
             with gr.Column():
                 out = gr.Textbox(lines=5, label="Output Text", value="")

dearth_config.py

@@ -46,7 +46,7 @@ class DearthConfig(PretrainedConfig):
         self.hidden_dim = hidden_dim
         if hidden_dim is None:
             self.hidden_dim = dim * 4
-            print(f"hidden_dim is not specified. Set to {self.hidden_dim}")
+            #print(f"hidden_dim is not specified. Set to {self.hidden_dim}")
         self.multiple_of = multiple_of
         self.dropout_rate = dropout_rate
         self.layer_init_factor = layer_init_factor
@@ -66,7 +66,7 @@ class DearthConfig(PretrainedConfig):
         self.mimic_use_alibi = mimic_use_alibi
 
         if "attn_window_size" in kwargs:
-            print("Warning: attn_window_size is deprecated. Please use sliding_window_size instead !!!!!!!!!!!")
+            #print("Warning: attn_window_size is deprecated. Please use sliding_window_size instead !!!!!!!!!!!")
             self.sliding_window_size = kwargs["attn_window_size"]
 
         super().__init__(

dearth_model.py

@@ -611,10 +611,10 @@ class DearthModel(nn.Module):
         self.residual_factor = config.residual_factor if config.residual_factor is not None else float(config.n_layer * 2) ** (1/4)
         if config.residual_factor is None:
             config.residual_factor = self.residual_factor
-            logging.warning(f"residual_factor is not set, using default value {self.residual_factor} = (2 * n_layer) ** 1/4")
+            #logging.warning(f"residual_factor is not set, using default value {self.residual_factor} = (2 * n_layer) ** 1/4")
         if config.layer_init_factor is None:
             config.layer_init_factor = self.layer_init_factor
-            logging.warning(f"layer_init_factor is not set, using default value {self.layer_init_factor} = (n_layer * 8) ** -1/2")
+            #logging.warning(f"layer_init_factor is not set, using default value {self.layer_init_factor} = (n_layer * 8) ** -1/2")
         
         self.config = config