Enhance visualization functionality and update data loading paths in Study_Models and Test_Model scripts

Enhance visualization process and model saving in Study_Models
Enhance file loading in Test_Model to include error handling and update validation pattern

Author: DefinetlyNotAI

.gitignore

@@ -208,3 +208,4 @@ __marimo__/
 /models/
 /cache/
 /data/
+/*Data_Visualization/

.idea/VulnScan.iml

@@ -286,8 +286,11 @@ def save_model_summary(model_, filename="Model_Summary.txt"):
 criterion = nn.BCEWithLogitsLoss()
 
 # Run all visualizations
+print("Running visualize_weight_distribution...")
 visualize_weight_distribution(model)
+print("Running visualize_activations...")
 visualize_activations(model, sample_input)
+print("Preparing texts and labels for t-SNE custom visualization...")
 texts = [
     # Non-sensitive (0)
     "I need to buy milk and bread from the grocery store.",
@@ -300,14 +303,23 @@ def save_model_summary(model_, filename="Model_Summary.txt"):
     "My social security number is 123-45-6789."
 ]
 labels = [0, 0, 0, 1, 1, 1]
+print("Loading SentenceTransformer embedder...")
 embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+print("Running visualize_tsne...")
 visualize_tsne(model, dataloader)
+print("Running visualize_tsne_custom...")
 visualize_tsne_custom(model, embedder, texts, labels)
+print("Running visualize_feature_importance...")
 visualize_feature_importance(input_dim)
+print("Running plot_loss_landscape_3d...")
 plot_loss_landscape_3d(model, dataloader, criterion)
+print("Saving model state dict...")
 save_model_state_dict(model)
+print("Generating model visualization...")
 generate_model_visualization(model, input_dim)
+print("Saving graph...")
 save_graph(model)
+print("Saving model summary...")
 save_model_summary(model)
 
 print("All visualizations completed. Files saved in 'data/' directory.")

tools/Study_Models.py

@@ -1,3 +1,5 @@
+import sys
+
 import torch
 from sentence_transformers import SentenceTransformer
 from vulnscan import SimpleNN
@@ -15,6 +17,9 @@
 def load_embeddings(folder_path, pattern):
     """Load all .pt files matching pattern and concatenate embeddings and labels"""
     files = sorted(glob.glob(os.path.join(folder_path, pattern)))
+    print("Found files:", files)
+    if not files:
+        sys.exit(f"No files found in {folder_path} matching {pattern}")
     all_embeddings = []
     all_labels = []
     for f in files:
@@ -27,18 +32,18 @@ def load_embeddings(folder_path, pattern):
 
 
 # Example paths
-cache_dir = f"cache/{NAME}/round_{ROUND}/embeddings"
+cache_dir = f"../cache/{NAME}/round_{ROUND}/embeddings"
 
 # Load all train/test/val embeddings
 train_embeddings, train_labels = load_embeddings(cache_dir, "train_*.pt")
 test_embeddings, test_labels = load_embeddings(cache_dir, "test_*.pt")
-val_embeddings, val_labels = load_embeddings(cache_dir, "val_*.pt")
+val_embeddings, val_labels = load_embeddings(cache_dir, "validation_*.pt")
 
 # Initialize model
 input_dim = train_embeddings.shape[1]
 model = SimpleNN(input_dim=input_dim).to(device)
 model.load_state_dict(torch.load(
-    f"cache/{NAME}/round_{ROUND}/{NAME}_round{ROUND}.pth",
+    f"../cache/{NAME}/round_{ROUND}/{NAME}_round{ROUND}.pth",
     map_location="cpu"
 ))
 model.eval()