add fct to plot confusion matrix

src/utils.py

@@ -1,4 +1,7 @@
+import matplotlib.pyplot as plt
+import seaborn as sns
 from sklearn.metrics import accuracy_score, f1_score
+from sklearn.metrics import confusion_matrix
 
 
 def get_metrics(y_true, y_preds):
@@ -16,4 +19,16 @@ def evaluate_predictions(model:str, train_preds, y_train, test_preds, y_test):
     get_metrics(y_train, train_preds)
     print("-"*50)
     print("Test set:")
-    get_metrics(y_test, test_preds)
+    get_metrics(y_test, test_preds)
+
+
+def plot_confusion_matrix(y_true, y_preds):
+    labels = sorted(set(y_true.tolist() + y_preds.tolist()))
+    cm = confusion_matrix(y_true, y_preds)
+    plt.figure(figsize=(12, 10))
+    sns.heatmap(cm, annot=True, cmap="Blues",
+                xticklabels=labels, yticklabels=labels)
+    plt.xlabel('Predicted Label')
+    plt.ylabel('True Label')
+    plt.title('Confusion Matrix')
+    plt.show()