mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
101 lines
2.7 KiB
Python
101 lines
2.7 KiB
Python
import os
|
|
import json
|
|
from flask import Flask, render_template, request, abort
|
|
|
|
app = Flask(__name__)
|
|
|
|
TRAIN_OUTPUTS_DIR = "train_outputs"
|
|
|
|
|
|
def get_run_data(run_path):
|
|
"""
|
|
Loads and processes data from all_results.json for a given run.
|
|
|
|
Args:
|
|
run_path: Path to the directory containing the all_results.json file.
|
|
|
|
Returns:
|
|
A dictionary containing the grouped results data.
|
|
"""
|
|
results_file = os.path.join(run_path, "all_results.json")
|
|
try:
|
|
with open(results_file, "r") as f:
|
|
data = json.load(f)
|
|
except FileNotFoundError:
|
|
return None
|
|
|
|
grouped_data = {}
|
|
for key, value in data.items():
|
|
if isinstance(value, float):
|
|
value = round(value, 4)
|
|
prefix = key.split("_")[0]
|
|
if prefix not in grouped_data:
|
|
grouped_data[prefix] = {}
|
|
grouped_data[prefix][key] = value
|
|
|
|
return grouped_data
|
|
|
|
|
|
def get_generated_text_data(run_path):
|
|
"""
|
|
Loads generated text data from test_generated_text.jsonl and val_generated_text.jsonl.
|
|
|
|
Args:
|
|
run_path: Path to the directory containing the .jsonl files.
|
|
|
|
Returns:
|
|
A dictionary containing the generated text data for "test" and "val".
|
|
"""
|
|
generated_data = {}
|
|
for split in ["test", "val"]:
|
|
filename = f"{split}_generated_text.jsonl"
|
|
filepath = os.path.join(run_path, filename)
|
|
try:
|
|
with open(filepath, "r") as f:
|
|
lines = f.readlines()
|
|
data = [json.loads(line) for line in lines]
|
|
generated_data[split] = data
|
|
except FileNotFoundError:
|
|
generated_data[split] = None
|
|
return generated_data
|
|
|
|
|
|
@app.route("/")
|
|
def index():
|
|
"""
|
|
Displays a dropdown list of training runs.
|
|
"""
|
|
runs = [
|
|
d
|
|
for d in os.listdir(TRAIN_OUTPUTS_DIR)
|
|
if os.path.isdir(os.path.join(TRAIN_OUTPUTS_DIR, d))
|
|
]
|
|
runs.sort(
|
|
key=lambda d: os.path.getctime(os.path.join(TRAIN_OUTPUTS_DIR, d)),
|
|
reverse=True,
|
|
)
|
|
return render_template("index.html", runs=runs)
|
|
|
|
|
|
@app.route("/visualize/<run>")
|
|
def visualize(run):
|
|
"""
|
|
Displays the results from all_results.json and the generated text data.
|
|
"""
|
|
logdir = os.path.join(TRAIN_OUTPUTS_DIR, run)
|
|
if not os.path.isdir(logdir):
|
|
abort(404, description=f"Run '{run}' not found.")
|
|
|
|
data = get_run_data(logdir)
|
|
if data is None:
|
|
abort(404, description=f"'all_results.json' not found in '{run}'.")
|
|
|
|
generated_data = get_generated_text_data(logdir)
|
|
|
|
return render_template(
|
|
"visualize.html", run=run, data=data, generated_data=generated_data
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run(debug=True)
|