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

import type { EChartsOption } from 'echarts'
import type { FineTuneTask, TrainingLogFile } from '@/types'
import type { FineTuneMetricPoint } from '@/api/modules/fineTune'

export interface TrainingMetricData {
  steps: number[]
  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 = { steps: [], loss: [], gradNorm: [], lr: [], epoch: [] }
  const candidates = text
    .split(/\r?\n/)
    .flatMap((line) => {
      const blocks = line.match(/\{[^{}\r\n]*\}/g)
      return blocks?.length ? blocks.map((block) => `${line} ${block}`) : [line]
    })

  for (const [index, line] of candidates.entries()) {
    const loss = extractNumber(line, 'loss')
    const gradNorm = extractNumber(line, 'grad_norm')
    const learningRate = extractNumber(line, 'learning_rate')
    const epoch = extractNumber(line, 'epoch')
    if (loss == null && gradNorm == null && learningRate == null) continue
    metrics.steps.push(extractNumber(line, 'step') ?? metrics.steps.length + index + 1)
    metrics.loss.push(loss ?? Number.NaN)
    metrics.gradNorm.push(gradNorm ?? Number.NaN)
    metrics.lr.push(learningRate ?? Number.NaN)
    metrics.epoch.push(epoch ?? Number.NaN)
  }

  return metrics
}

export function metricsFromApi(points: FineTuneMetricPoint[]): TrainingMetricData {
  const metrics: TrainingMetricData = { steps: [], loss: [], gradNorm: [], lr: [], epoch: [] }
  for (const [index, point] of points.entries()) {
    const hasMetric = point.loss != null || point.grad_norm != null || point.learning_rate != null
    if (!hasMetric) continue
    metrics.steps.push(Number(point.step || index + 1))
    metrics.loss.push(point.loss == null ? Number.NaN : Number(point.loss))
    metrics.gradNorm.push(point.grad_norm == null ? Number.NaN : Number(point.grad_norm))
    metrics.lr.push(point.learning_rate == null ? Number.NaN : Number(point.learning_rate))
    metrics.epoch.push(point.epoch == null ? Number.NaN : Number(point.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[],
  steps: number[],
  color: string,
  logScale = false,
): EChartsOption {
  const visibleData = data.map((value) => (Number.isFinite(value) ? value : null))
  return {
    grid: { top: 24, right: 20, bottom: 56, left: 56 },
    graphic: visibleData.some((value) => value != null)
      ? []
      : [
          {
            type: 'text',
            left: 'center',
            top: 'middle',
            style: { text: '暂无训练指标数据', fill: '#94a3b8', fontSize: 13 },
          },
        ],
    tooltip: {
      trigger: 'axis',
      axisPointer: { type: 'cross' },
      backgroundColor: 'rgba(15, 23, 42, 0.9)',
      borderWidth: 0,
      textStyle: { color: '#fff', fontSize: 12 },
    },
    xAxis: {
      type: 'category',
      data: steps.map((step, index) => (Number.isFinite(step) ? String(step) : String(index + 1))),
      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: visibleData,
        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` },
            ],
          },
        },
      },
    ],
  }
}