YG_FT/frontend/src/views/system/training-log/trainingLogModel.ts

import type { EChartsOption } from 'echarts'
import type { FineTuneTask, TrainingLogFile } from '@/types'

export interface TrainingMetricData {
  loss: number[]
  gradNorm: number[]
  lr: number[]
  epoch: number[]
}

export interface TrainingSummary {
  epoch: string
  trainLoss: string
  runtime: string
}

export interface ParsedTrainingLog {
  metrics: TrainingMetricData
  summary: TrainingSummary
}

const NUMBER_SOURCE = '[-+]?(?:\\d+(?:\\.\\d*)?|\\.\\d+)(?:[eE][-+]?\\d+)?'

function escapeRegExp(value: string) {
  return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')
}

function extractNumber(source: string, key: string) {
  const match = source.match(new RegExp(`['"]?${escapeRegExp(key)}['"]?\\s*:\\s*(${NUMBER_SOURCE})`, 'i'))
  return match ? Number(match[1]) : undefined
}

function extractSummaryValue(source: string, key: string) {
  const match = source.match(new RegExp(`['"]?${escapeRegExp(key)}['"]?\\s*(?:=|:)\\s*(${NUMBER_SOURCE})`, 'i'))
  return match?.[1] || ''
}

/** 根据任务精确选择日志；PID 优先，任务名仅作为明确兜底。 */
export function resolveTrainingLogFile(
  files: TrainingLogFile[],
  task: Pick<FineTuneTask, 'process_id' | 'name'>,
) {
  const processId = task.process_id
  if (processId != null) {
    const pidMatch = files.find((file) => file.pid === processId)
    if (pidMatch) return pidMatch

    const pidPattern = new RegExp(`(?:^|[^0-9])(?:pid)?${processId}(?:[^0-9]|$)`, 'i')
    const filenameMatch = files.find((file) => pidPattern.test(file.file))
    if (filenameMatch) return filenameMatch
  }

  const taskName = task.name.trim()
  if (!taskName) return undefined
  return files.find((file) => file.name.includes(taskName) || file.file.includes(taskName))
}

/** 解析日志中的逐步指标。字段顺序和常见数值格式均不受限制。 */
export function parseTrainingMetrics(text: string): TrainingMetricData {
  const metrics: TrainingMetricData = { loss: [], gradNorm: [], lr: [], epoch: [] }
  const blocks = text.match(/\{[^{}\r\n]*\}/g) || []

  for (const block of blocks) {
    const loss = extractNumber(block, 'loss')
    const gradNorm = extractNumber(block, 'grad_norm')
    const learningRate = extractNumber(block, 'learning_rate')
    const epoch = extractNumber(block, 'epoch')
    if (loss == null || gradNorm == null || learningRate == null) continue
    metrics.loss.push(loss)
    metrics.gradNorm.push(gradNorm)
    metrics.lr.push(learningRate)
    if (epoch != null) metrics.epoch.push(epoch)
  }

  return metrics
}

/** 每次都返回新对象，日志截断或切换时不会残留上一轮汇总。 */
export function parseTrainingSummary(text: string): TrainingSummary {
  const emptySummary: TrainingSummary = { epoch: '', trainLoss: '', runtime: '' }
  const startMatch = /\*{5}\s*train metrics\s*\*{5}/i.exec(text)
  if (!startMatch) return emptySummary

  const tail = text.slice(startMatch.index + startMatch[0].length)
  const endMatch = /\*{5}\s*train metrics end\s*\*{5}/i.exec(tail)
  const body = endMatch ? tail.slice(0, endMatch.index) : tail
  return {
    epoch: extractSummaryValue(body, 'epoch'),
    trainLoss: extractSummaryValue(body, 'train_loss'),
    runtime: extractSummaryValue(body, 'train_runtime'),
  }
}

export function parseTrainingLog(text: string): ParsedTrainingLog {
  return {
    metrics: parseTrainingMetrics(text),
    summary: parseTrainingSummary(text),
  }
}

/** 构建单条训练指标曲线。 */
export function buildMetricChartOption(
  label: string,
  data: number[],
  color: string,
  logScale = false,
): EChartsOption {
  return {
    grid: { top: 24, right: 20, bottom: 56, left: 56 },
    tooltip: {
      trigger: 'axis',
      axisPointer: { type: 'cross' },
      backgroundColor: 'rgba(15, 23, 42, 0.9)',
      borderWidth: 0,
      textStyle: { color: '#fff', fontSize: 12 },
    },
    xAxis: {
      type: 'category',
      boundaryGap: false,
      name: 'Step',
      nameTextStyle: { color: '#94a3b8', fontSize: 11 },
      axisLine: { lineStyle: { color: '#e2e8f0' } },
      axisLabel: { color: '#94a3b8', fontSize: 11 },
      splitLine: { show: false },
    },
    yAxis: {
      type: logScale ? 'log' : 'value',
      name: label,
      nameTextStyle: { color: '#94a3b8', fontSize: 11 },
      axisLine: { show: false },
      axisTick: { show: false },
      axisLabel: { color: '#94a3b8', fontSize: 11 },
      splitLine: { lineStyle: { color: '#f1f5f9' } },
    },
    dataZoom: data.length > 30
      ? [
          { type: 'inside', start: 0, end: 100 },
          { type: 'slider', height: 16, bottom: 8, borderColor: 'transparent', fillerColor: 'rgba(79,70,229,0.08)', handleStyle: { color: '#4f46e5' } },
        ]
      : [],
    series: [
      {
        name: label,
        type: 'line',
        data,
        smooth: true,
        symbol: 'none',
        lineStyle: { width: 2, color },
        areaStyle: {
          color: {
            type: 'linear',
            x: 0,
            y: 0,
            x2: 0,
            y2: 1,
            colorStops: [
              { offset: 0, color: `${color}55` },
              { offset: 1, color: `${color}05` },
            ],
          },
        },
      },
    ],
  }
}