import wandb
import re

# Initialize WandB project
wandb.init(project="training-amharic-stt-visualizations", name="metrics_visualization")

# Path to your training log file
log_file = "Training/train_log.txt"

# Function to parse logs
# Function to parse logs
def parse_logs(log_file):
    """
    Parses the training logs and yields metrics as dictionaries.
    """
    with open(log_file, "r") as f:
        for line in f:
            # Match the log format using regex
            match = re.match(
                r"epoch: (?P<epoch>\d+), lr_model: (?P<lr_model>[0-9.e+-]+), lr_wav2vec: (?P<lr_wav2vec>[0-9.e+-]+) - "
                r"train loss: (?P<train_loss>[0-9.e+-]+) - valid loss: (?P<valid_loss>[0-9.e+-]+), "
                r"valid CER: (?P<valid_CER>[0-9.e+-]+), valid WER: (?P<valid_WER>[0-9.e+-]+)",
                line.strip()
            )
            if match:
                metrics = {key: float(value) if '.' in value or 'e' in value else int(value)
                           for key, value in match.groupdict().items()}
                yield metrics

# Parse logs and log to WandB
for metrics in parse_logs(log_file):
    # Log metrics to WandB
    wandb.log(metrics)

# Finish WandB run
wandb.finish()