diff --git a/webui/app.py b/webui/app.py
index 0985cbb..230b623 100644
--- a/webui/app.py
+++ b/webui/app.py
@@ -1,6 +1,7 @@
 import json
 import os
 import re
+from glob import glob
 
 import torch
 import yaml
@@ -47,20 +48,23 @@ def get_run_data(run_path):
 
 def get_generated_text_data(run_path):
     """
-    Loads generated text data from test_generated_text.jsonl and val_generated_text.jsonl.
+    Loads generated text data from all *_generated_text.jsonl files.
 
     Args:
         run_path: Path to the directory containing the .jsonl files.
 
     Returns:
-        A dictionary containing the generated text data for "test" and "val".
+        A dictionary containing the generated text data for each split found.
     """
     generated_data = {}
-    for split in ["test", "val"]:
-        filename = f"{split}_generated_text.jsonl"
-        filepath = os.path.join(run_path, filename)
+    files = sorted(glob(f"{run_path}/*_generated_text.jsonl"))
+    print(files)
+    # Find all *_generated_text.jsonl files
+    for filename in files:
+        # Extract split name from filename (remove _generated_text.jsonl)
+        split = filename.split("/")[-1].split("_generated_text.jsonl")[0]
         try:
-            with open(filepath) as f:
+            with open(filename) as f:
                 lines = f.readlines()
                 data = [json.loads(line) for line in lines]
             generated_data[split] = data
diff --git a/webui/templates/visualize.html b/webui/templates/visualize.html
index be0934a..c327771 100644
--- a/webui/templates/visualize.html
+++ b/webui/templates/visualize.html
@@ -249,6 +249,7 @@
     </div>
 
     <div class="column">
+        {% if generated_data %}
         {% for split, data in generated_data.items() %}
         {% if data %}
         <div class="generated-text-container">
@@ -260,7 +261,6 @@
                 <p><strong>Input:</strong> <span id="{{ split }}-input"></span></p>
                 <p><strong>Generated:</strong> <span id="{{ split }}-generated"></span></p>
                 <p><strong>Label:</strong> <span id="{{ split }}-label"></span></p>
-
             </div>
             <div class="generated-text-controls">
                 <button onclick="prev('{{ split }}')">Prev</button>
@@ -273,6 +273,9 @@
         </div>
         {% endif %}
         {% endfor %}
+        {% else %}
+        <p>No generated text data found.</p>
+        {% endif %}
 
         <div id="load-chat-container">
             <button id="load-chat-button">Load Chat with {{ model_name }}</button>