# ==============================================================================
# COPYRIGHT (C) 2025 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
#
# This software is licensed under the Commercial License Agreement V.1.2.
# Any use, modification, or distribution of this code requires compliance with 
# the terms found in the LICENSE.md file in the root directory.
#
# NO PATENTING RIGHTS: Users are strictly prohibited from filing patent claims
# based on the BRE or SWA architectures disclosed herein.
# Contact: grabko@cmsmanhattan.com | +1 (516) 777-0945
# ==============================================================================


import matplotlib.pyplot as plt
import re
import os

LOG_FILE = "fine_tune_pile.log"
OUTPUT_IMAGE = "loss_chart.png"

def plot_loss():
    if not os.path.exists(LOG_FILE):
        print(f"Файл {LOG_FILE} не найден.")
        return

    steps = []
    losses = []

    # Регулярное выражение для поиска шага и лосса
    pattern = re.compile(r"Шаг (\d+).*Loss: (\d+\.\d+)")

    with open(LOG_FILE, "r") as f:
        for line in f:
            match = pattern.search(line)
            if match:
                steps.append(int(match.group(1)))
                losses.append(float(match.group(2)))

    if not steps:
        print("Данные для графика пока не найдены в логе.")
        return

    plt.figure(figsize=(10, 6))
    plt.plot(steps, losses, label='Raw Loss', alpha=0.3, color='blue')
    
    # Скользящее среднее для сглаживания
    if len(losses) > 10:
        window = 10
        smooth_loss = [sum(losses[i:i+window])/window for i in range(len(losses)-window)]
        plt.plot(steps[window:], smooth_loss, label='Smooth Loss (EMA)', color='red', linewidth=2)

    plt.xlabel('Step')
    plt.ylabel('Loss')
    plt.title('JiRack 1B Training Progress')
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.6)
    
    plt.savefig(OUTPUT_IMAGE)
    print(f"График сохранен в {OUTPUT_IMAGE}")

if __name__ == "__main__":
    plot_loss()