import re
import wandb

wandb.init(project="gpt2_training", name="xs_pretrain_small_visualization7")

log_file = "temp_log7/xs_pretrain_small_7.txt"

try:
    with open(log_file, "r", encoding="utf-8") as f:
        log_lines = f.read()
except FileNotFoundError:
    print(f"Error: File {log_file} not found!")
    exit(1)

train_steps = []
train_losses = []
val_steps = []
val_losses = []

for line_num, line in enumerate(log_lines.splitlines(), 1):
    if not line.strip(): 
        continue
    train_match = re.match(r"step:(\d+)/\d+ train_loss:([\d.]+) train_time:(nan|\d+)ms step_avg:(nan|\d+\.\d+)ms", line)
    if train_match:
        step, train_loss = int(train_match.group(1)), float(train_match.group(2))
        train_steps.append(step)
        train_losses.append(train_loss)
        wandb.log({"step": step, "train_loss": train_loss})
    else:
        val_match = re.match(r"step:(\d+)/\d+ val_loss:([\d.]+) train_time:(nan|\d+)ms step_avg:(nan|\d+\.\d+)ms", line)
        if val_match:
            step, val_loss = int(val_match.group(1)), float(val_match.group(2))
            val_steps.append(step)
            val_losses.append(val_loss)
            wandb.log({"step": step, "val_loss": val_loss})
        else:
            print(f"Line {line_num} failed to match: '{line}'")
if not train_steps:
    print("Error: No training loss data found in the log file!")
    wandb.finish()
    exit(1)
if not val_steps:
    print("Warning: No validation loss data found in the log file!")
wandb.finish()
print("Visualization completed in WandB! Check your dashboard at wandb.ai")
