Update app.py

app.py

@@ -7,6 +7,7 @@ from transformers import (
 )
 from peft import PeftModel
 import torch
+import time
 
 model_path = "Qwen/Qwen1.5-1.8B-Chat"
 lora_path = "AngoHF/EssayGPT" #+ "/checkpoint-100"
@@ -31,11 +32,7 @@ model = PeftModel.from_pretrained(model, lora_path)
 model = model.merge_and_unload()
 model.eval()
 
-# model.config.use_cache = True
-# model.to("cpu")
-# model.save_pretrained("/data/ango/EssayGPT")
-
-# tokenizer.save_pretrained("/data/ango/EssayGPT")
+model.config.use_cache = True
 
 
 MAX_MATERIALS = 4
@@ -55,11 +52,14 @@ def call(related_materials, materials, question):
         add_generation_prompt=True
     )
     model_inputs = tokenizer([text], return_tensors="pt").to(device)
-    print(len(model_inputs.input_ids[0]))
+    print(f"Input Token Length: {len(model_inputs.input_ids[0])}")
+    start_time = time.time()
     generated_ids = model.generate(
         model_inputs.input_ids,
         max_length=8096
     )
+    
+    print(f"Inference Cost Time: {time.time() - start_time}")
     generated_ids = [
         output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
     ]
@@ -107,7 +107,7 @@ def build_ui(components):
 def run():
     app = create_ui()
     app.queue()
-    app.launch(share=True)
+    app.launch()
 
 
 if __name__ == '__main__':