add time

app.py

@@ -149,6 +149,8 @@ def predict(
     model_inputs = tokenizer([text_for_model], return_tensors="pt").to(model.device)
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     # print("start generating")
+
+    loc_start_time = time.time()
     if temperature > 0:
         generation_kwargs = dict(
             **model_inputs,
@@ -177,12 +179,17 @@ def predict(
         yield [[prompt, highlight_response]], [["null test", "null test2"]], "Generating file locations..."
     
     gen_thread.join()
+    loc_end_time = time.time()
+    loc_time = loc_end_time - loc_start_time
+
+    encoded_answer = tokenizer(partial_output, padding=True, truncation=True, return_tensors='pt')
+    print("loc token/s:",len(encoded_answer['input_ids'][0])/loc_time)
 
     response = partial_output
 
     raw_answer=post_process(response)
     model_found_files = raw_answer.strip().split("\n")
-    print(response)
+    # print(response)
 
     highlight_response = highlight_thinking(response)
     yield [[prompt,highlight_response]], [["null test","null test2"]], "Generate: Success"
@@ -194,7 +201,7 @@ def predict(
         print(file_name)
         # pdb.set_trace()
         to_open_path = repo_path + "/" + file_name
-        print("to_open_path,",to_open_path)
+
         with open(to_open_path, "r", encoding="utf-8") as f:
             content = f.read()
         contents += f"{file_name}\n{content}\n\n"
@@ -214,6 +221,7 @@ def predict(
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
     subprocess.run(["rm", "-rf", repo_path], check=True)
+    repair_start_time = time.time()
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     if temperature > 0:
         generation_kwargs = dict(
@@ -244,6 +252,12 @@ def predict(
         yield [[prompt, highlight_response], [repair_prompt, highlight_response_repair]], [["null test", "null test2"]], "Generating repair suggestion..."
 
     gen_thread.join()
+    repair_end_time = time.time()
+
+    repair_time = repair_end_time - repair_start_time
+
+    encoded_answer = tokenizer(partial_output_repair, padding=True, truncation=True, return_tensors='pt')
+    print("loc token/s:",len(encoded_answer['input_ids'][0])/repair_time)
 
     # yield response, "null test", "Generate: Success"
     yield [[prompt,highlight_response],[repair_prompt,highlight_response_repair]], [["null test","null test2"]], "Generate: Success"