fix slow respond after days

app.py

@@ -9,19 +9,9 @@ from dearth_model import DearthForCausalLM
 import random
 
 
-tk = transformers.AutoTokenizer.from_pretrained("./tk")
 
-#model_path = "./ts100-re2-h1-4000.pt"
+tk = transformers.AutoTokenizer.from_pretrained("./tk")
 model_path = "./ts100-re2-h1-4000-model.pt"
-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)
 states = torch.load(model_path, map_location="cpu")
 model_states = states
 unwanted_prefix_dueto_compile = '_orig_mod.'
@@ -39,10 +29,21 @@ for k,v in list(model_states.items()):
         new_key = k[len(unwanted_prefix_dueto_compile):]
         model_states[k[len(unwanted_prefix_dueto_compile):]] = model_states.pop(k)
 
-model.load_state_dict(model_states)
+def generate(input, num_more_tokens):
+  
+    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)
 
 
-def generate(input, num_more_tokens):
     num_more_tokens = int(num_more_tokens)
     print(input)
     input = input.strip()
@@ -84,7 +85,7 @@ The PPL on the validation set is 1.7, in comparison, the teacher model has a PPL
 """
 
 
-server = gr.Interface(
+demo = gr.Interface(
     fn=generate,
     title="Tinystories LM 11M",
     description=Description,
@@ -95,4 +96,6 @@ server = gr.Interface(
     outputs="text"
 )
 
-server.launch()
+if __name__ == "__main__":
+    demo.queue()
+    demo.launch(show_api=False)