Refactor. Extracted methods.

result_data_processor.py

@@ -4,65 +4,58 @@ import fnmatch
 import json
 
 class ResultDataProcessor:
-    def __init__(self):
+    
+    def __init__(self, directory='results', pattern='results*.json'):
+        self.directory = directory
+        self.pattern = pattern
         self.data = self.process_data()
 
+    @staticmethod
+    def _find_files(directory, pattern):
+        for root, dirs, files in os.walk(directory):
+            for basename in files:
+                if fnmatch.fnmatch(basename, pattern):
+                    filename = os.path.join(root, basename)
+                    yield filename
+    
+    def _read_and_transform_data(self, filename):
+        with open(filename) as f:
+            data = json.load(f)
+        df = pd.DataFrame(data['results']).T
+        return df
+    
+    def _cleanup_dataframe(self, df, model_name):
+        df = df.rename(columns={'acc': model_name})
+        df.index = (df.index.str.replace('hendrycksTest-', 'MMLU_', regex=True)
+                          .str.replace('harness\|', '', regex=True)
+                          .str.replace('\|5', '', regex=True))
+        return df[[model_name]]
+    
     def process_data(self):
-        dataframes = []
-
-        def find_files(directory, pattern):
-            for root, dirs, files in os.walk(directory):
-                for basename in files:
-                    if fnmatch.fnmatch(basename, pattern):
-                        filename = os.path.join(root, basename)
-                        yield filename
-
-        for filename in find_files('results', 'results*.json'):
-            model_name = filename.split('/')[2]
-            with open(filename) as f:
-                data = json.load(f)
-                df = pd.DataFrame(data['results']).T
-
-
-                # data cleanup
-                df = df.rename(columns={'acc': model_name})
-                # Replace 'hendrycksTest-' with a more descriptive column name
-                df.index = df.index.str.replace('hendrycksTest-', 'MMLU_', regex=True)
-                df.index = df.index.str.replace('harness\|', '', regex=True)
-                # remove |5 from the index
-                df.index = df.index.str.replace('\|5', '', regex=True)
+        dataframes = [self._cleanup_dataframe(self._read_and_transform_data(filename), filename.split('/')[2])
+                      for filename in self._find_files(self.directory, self.pattern)]
 
-
-                dataframes.append(df[[model_name]])
-
-        data = pd.concat(dataframes, axis=1)
-
-        data = data.transpose()
+        data = pd.concat(dataframes, axis=1).transpose()
+        
+        # Add Model Name and rearrange columns
         data['Model Name'] = data.index
         cols = data.columns.tolist()
         cols = cols[-1:] + cols[:-1]
         data = data[cols]
 
-        # remove the Model Name column
-        data = data.drop(['Model Name'], axis=1)
-
-        # remove the all column
-        data = data.drop(['all'], axis=1)
-
-        # remove the truthfulqa:mc|0 column
-        data = data.drop(['truthfulqa:mc|0'], axis=1)
-
-        # create a new column that averages the results from each of the columns with a name that start with MMLU
+        # Remove the 'Model Name' column
+        data = data.drop(columns=['Model Name'])
+        
+        # Add average column
         data['MMLU_average'] = data.filter(regex='MMLU').mean(axis=1)
 
-        # move the MMLU_average column to the third column in the dataframe
+        # Reorder columns to move 'MMLU_average' to the third position
         cols = data.columns.tolist()
         cols = cols[:2] + cols[-1:] + cols[2:-1]
         data = data[cols]
 
-        return data
-    
-    # filter data based on the index
+        # Drop specific columns
+        return data.drop(columns=['all', 'truthfulqa:mc|0'])
+
     def get_data(self, selected_models):
-        filtered_data = self.data[self.data.index.isin(selected_models)]
-        return filtered_data
+        return self.data[self.data.index.isin(selected_models)]