plot training history

src/model_training.py

@@ -1,6 +1,3 @@
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -11,7 +8,7 @@ from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-from utils import get_dataset
+from utils import get_dataset, plot_training_history
 
 
 def train_model(model, optimizer, train_loader, val_loader, num_epochs=100, patience=10):
@@ -125,7 +122,8 @@ def main():
     num_epochs = 100
 
     model, history = train_model(model, optimizer, train_loader, val_loader, num_epochs=num_epochs)
-        
+    plot_training_history(history)
+    
 
 if __name__ == "__main__":
     main()

src/utils.py

@@ -1,6 +1,7 @@
 import pickle
 
 import matplotlib.pyplot as plt
+import numpy as np
 import pandas as pd
 import seaborn as sns
 from datasets import DatasetDict, Dataset
@@ -68,4 +69,28 @@ def plot_confusion_matrix(y_true, y_preds):
     plt.xlabel('Predicted Label')
     plt.ylabel('True Label')
     plt.title('Confusion Matrix')
-    plt.show()
+    plt.show()
+
+
+def plot_training_history(history):
+    epochs = np.arange(1, len(history["train_loss"]) + 1)
+    plt.figure(figsize=(10, 5))
+
+    plt.subplot(1, 2, 1)
+    plt.plot(epochs, history["train_loss"], label='Train Loss')
+    plt.plot(epochs, history["valid_loss"], label='Valid Loss')
+    plt.xlabel('Epoch')
+    plt.ylabel('Loss')
+    plt.title('Training and Validation Loss')
+    plt.legend()
+
+    plt.subplot(1, 2, 2)
+    plt.plot(epochs, history["train_accuracies"], label='Train Accuracy')
+    plt.plot(epochs, history["valid_accuracies"], label='Valid Accuracy')
+    plt.xlabel('Epoch')
+    plt.ylabel('Accuracy')
+    plt.title('Training and Validation Accuracy')
+    plt.legend()
+
+    plt.tight_layout()
+    plt.savefig('../docs/images/training_history.png')