import { TensorStatistics } from '../checks/weights';
import { LoRAConfig } from '../models/config';
import { NamedTensorMap, Tensor } from '@tensorflow/tfjs-core';
export interface Metrics {
    accuracy?: number;
    perplexity?: number;
    loss: number;
}
export interface TrainingLogEntry {
    trainingMetrics: Metrics;
    validationMetrics?: Metrics;
    step: number;
    time: number;
    example?: string;
    batchSize: number;
    gradientMetrics?: Map<string, TensorStatistics>;
    learningRate?: number;
    gradientNorm?: number;
    weightNorm?: number;
    weightStatistics?: Map<string, TensorStatistics>;
    memoryUsage?: number;
    tokensPerSecond: number;
    duration: number;
    totalTokens: number;
}
export interface LRSchedulerConfig {
    warmupSteps: number;
    decayEpochs: number;
    minLearningRate: number;
    epochSteps: number;
    step?: number;
}
export interface AdamWOptimizerConfig extends LRSchedulerConfig {
    learningRate: number;
    beta1: number;
    beta2: number;
    accBeta1?: number;
    accBeta2?: number;
    epsilon?: number;
    weightDecay: number;
    lossScaling: number;
    clipNorm?: number;
    orthoGrad?: boolean;
}
export type TrainingMetrics = 'accuracy' | 'perplexity' | 'gradientNorm' | 'gradientStatistics' | 'weightNorm' | 'weightStatistics' | 'memoryUsage' | 'tokensPerSecond' | 'learningRate';
export interface TrainingOptions extends Partial<AdamWOptimizerConfig> {
    batchSize: number;
    maxEpochs?: number;
    logInterval?: number;
    prompt?: string;
    validationSplit?: number;
    gradientCheckpointing?: boolean;
    mixedPrecision?: boolean;
    trainableWeights?: string[];
    loraConfig?: LoRAConfig;
    loraName?: string;
    sftMode: 'full' | 'lora' | 'last-layer';
    maskedLoss?: boolean;
    metrics?: TrainingMetrics[];
    contextScaling?: number;
    labelSmoothing?: number;
    dropout?: number;
    layerDrop?: number;
    onStep?: (log: TrainingLogEntry) => void;
}
export interface TrainingState {
    step: number;
    lastLoss: number;
    totalSteps: number;
    losses: number[];
    validationLosses: number[];
    logStartTime: number;
    trainingDuration: number;
    gradients?: NamedTensorMap;
    gradientNorm?: Tensor;
    accuracy?: Tensor;
}
