use json built in module and os.listdir

extract.py

@@ -1,10 +1,8 @@
-import subprocess
 import re
-import matplotlib.pyplot as plt
-import datetime
+import json
 import numpy as np
 import statistics
-import pdb
+import os
 import csv
 
 model = []
@@ -18,41 +16,23 @@ temp_response_length = []
 temp_latency = []
 temp_energy = []
 
-model1 = input("model 1: ")
-model2 = input("model 2: ")
-model3 = input("model 3: ")
-model4 = input("model 4: ")
-
-model_name = []
-model_name.append(model1)
-model_name.append(model2)
-model_name.append(model3)
-model_name.append(model4)
+model_name = os.listdir("data/chat")
 
 match_name = False
 
-for i in range(len(model_name)):
-    with open(model_name[i], 'r') as file:
-        model_lines = file.readlines()
+for models in model_name:
+    with open("data/chat/"+models+"/benchmark.json", 'r') as file:
+        json_data = json.load(file)
 
-    for i in range(len(model_lines)):
-        match = re.search(r'"model":\s*"([^"]+)"', model_lines[i])
-        match1 = re.search(r'"throughput":\s*(\d+.\d+)', model_lines[i])
-        match2 = re.search(r'"response_length":\s*([0-9]+)', model_lines[i])
-        match3 = re.search(r'"latency":\s*(\d+.\d+)', model_lines[i])
-        match4 = re.search(r'"energy":\s*(\d+.\d+)', model_lines[i])
-        if match and not match_name:
-            temp_model_name = str(match.group(1))
-            model.append(temp_model_name.replace('--', '/'))
+    for obj in json_data:
+        if not match_name:
+            name = str(obj["model"])
+            model.append(name.replace('--','/'))
             match_name = True
-        elif match1:
-            temp_throughput.append(float(match1.group(1)))
-        elif match2:
-            temp_response_length.append(float(match2.group(1)))
-        elif match3:
-            temp_latency.append(float(match3.group(1)))
-        elif match4:
-            temp_energy.append(float(match4.group(1)))
+        temp_throughput.append(float(obj["throughput"]))
+        temp_response_length.append(float(obj["response_length"]))
+        temp_latency.append(float(obj["latency"]))
+        temp_energy.append(float(obj["energy"]))
         
     match_name = False
 
@@ -67,10 +47,10 @@ for i in range(len(model_name)):
     temp_energy.clear()
 
 
-avg_throughput = [sum(row) / len(row) for row in throughput]
-avg_response_length = [sum(row) / len(row) for row in response_length]
-avg_latency = [sum(row) / len(row) for row in latency]
-avg_energy = [sum(row) / len(row) for row in energy]
+avg_throughput = [statistics.mean(row) for row in throughput]
+avg_response_length = [statistics.mean(row) for row in response_length]
+avg_latency = [statistics.mean(row) for row in latency]
+avg_energy = [statistics.mean(row) for row in energy]
 
 for i in range(len(model)):
     print(model[i])