condensed functionality in a single function for plotting both epoch and step metrics

examples/plot_metrics.py

@@ -1,69 +1,53 @@
 import json
 import matplotlib.pyplot as plt
-import sys
+import argparse
 import os
 
-def plot_metrics(file_path):
-
-    # Read the JSON file
-    with open(file_path, 'r') as f:
-        data = json.load(f)
-
-    # Get directory and filename information
-    directory = os.path.dirname(file_path)
-    filename_prefix = os.path.basename(file_path).split('.')[0]
-
-    # Plotting metrics for training and validation step loss
+def plot_metric(data, metric_name, x_label, y_label, title, colors):
     plt.figure(figsize=(14, 6))
     plt.subplot(1, 2, 1)
-    plt.plot(data['train_step_loss'], label='Train Step Loss', color='b')
-    plt.plot(data['val_step_loss'], label='Validation Step Loss', color='r')
-    plt.xlabel('Step')
-    plt.ylabel('Loss')
-    plt.title('Train and Validation Step Loss')
+    plt.plot(data[f'train_step_{metric_name}'], label=f'Train Step {metric_name.capitalize()}', color=colors[0])
+    plt.plot(data[f'val_step_{metric_name}'], label=f'Validation Step {metric_name.capitalize()}', color=colors[1])
+    plt.xlabel(x_label)
+    plt.ylabel(y_label)
+    plt.title(f'Train and Validation Step {title}')
     plt.legend()
 
-    # Plotting metrics for training and validation epoch loss
     plt.subplot(1, 2, 2)
-    plt.plot(data['train_epoch_loss'], label='Train Epoch Loss', color='b')
-    plt.plot(data['val_epoch_loss'], label='Validation Epoch Loss', color='r')
-    plt.xlabel('Epoch')
-    plt.ylabel('Loss')
-    plt.title('Train and Validation Epoch Loss')
+    plt.plot(data[f'train_epoch_{metric_name}'], label=f'Train Epoch {metric_name.capitalize()}', color=colors[0])
+    plt.plot(data[f'val_epoch_{metric_name}'], label=f'Validation Epoch {metric_name.capitalize()}', color=colors[1])
+    plt.xlabel(x_label)
+    plt.ylabel(y_label)
+    plt.title(f'Train and Validation Epoch {title}')
     plt.legend()
     plt.tight_layout()
+
+def plot_metrics(file_path):
+    if not os.path.exists(file_path):
+        print(f"File {file_path} does not exist.")
+        return
+
+    with open(file_path, 'r') as f:
+        try:
+            data = json.load(f)
+        except json.JSONDecodeError:
+            print("Invalid JSON file.")
+            return
+
+    directory = os.path.dirname(file_path)
+    filename_prefix = os.path.basename(file_path).split('.')[0]
+
+    plot_metric(data, 'loss', 'Step', 'Loss', 'Loss', ['b', 'r'])
     plt.savefig(os.path.join(directory, f"{filename_prefix}_train_and_validation_loss.png"))
     plt.close()
 
-    # Plotting perplexity
-    plt.figure(figsize=(14, 6))
-    plt.subplot(1, 2, 1)
-    plt.plot(data['train_step_perplexity'],
-             label='Train Step Perplexity', color='g')
-    plt.plot(data['val_step_perplexity'],
-             label='Validation Step Perplexity', color='m')
-    plt.xlabel('Step')
-    plt.ylabel('Perplexity')
-    plt.title('Train and Validation Step Perplexity')
-    plt.legend()
-    
-    plt.subplot(1, 2, 2)
-    plt.plot(data['train_epoch_perplexity'],
-             label='Train Epoch Perplexity', color='g')
-    plt.plot(data['val_epoch_perplexity'],
-             label='Validation Epoch Perplexity', color='m')
-    plt.xlabel('Epoch')
-    plt.ylabel('Perplexity')
-    plt.title('Train and Validation Epoch Perplexity')
-    plt.legend()
-    plt.tight_layout()
+    plot_metric(data, 'perplexity', 'Step', 'Perplexity', 'Perplexity', ['g', 'm'])
     plt.savefig(os.path.join(directory, f"{filename_prefix}_train_and_validation_perplexity.png"))
     plt.close()
 
 if __name__ == "__main__":
-    if len(sys.argv) < 2:
-        print("Usage: python script.py <path_to_metrics_json>")
-        sys.exit(1)
+    parser = argparse.ArgumentParser(description='Plot metrics from JSON file.')
+    parser.add_argument('file_path', type=str, help='Path to the metrics JSON file.')
+    args = parser.parse_args()
 
-    file_path = sys.argv[1]
-    plot_metrics(file_path)
+    plot_metrics(args.file_path)