import { ITokeniser } from '../tokeniser/type';
import { Scalar, Tensor } from '@tensorflow/tfjs-core';
import { Dataset } from '@tensorflow/tfjs-data';
import { default as Model, ModelForwardAttributes } from '../models/model';
import { AdamWOptimizerConfig, TrainingLogEntry, TrainingMetrics, TrainingOptions, TrainingState } from './types';
import { AdamWOptimizer } from './AdamW';
export default class BasicTrainer {
    tokenizer: ITokeniser;
    model: Model<ModelForwardAttributes>;
    optimizer: AdamWOptimizer;
    protected running: boolean;
    protected lastState?: TrainingState;
    protected _gradientCheckpointing: boolean;
    protected _mixedPrecision: boolean;
    protected maskedLoss: boolean;
    protected optimizerConfig: AdamWOptimizerConfig;
    protected metrics: Set<TrainingMetrics>;
    protected _labelSmoothing: number;
    protected _layerDrop: number;
    protected _dropout: number;
    constructor(model: Model<ModelForwardAttributes>, tokenizer: ITokeniser, optConfig?: Partial<AdamWOptimizerConfig>, optimizer?: AdamWOptimizer);
    setGradientCheckpointing(enabled: boolean): void;
    setMixedPrecision(enabled: boolean): void;
    setLabelSmoothing(smoothing: number): void;
    setDropout(dropout: number): void;
    setLayerDrop(layerDrop: number): void;
    setLearningRate(learningRate: number): void;
    setMetrics(metrics: TrainingMetrics[]): void;
    reset(): void;
    stop(): void;
    get isRunning(): boolean;
    getOptimizer(): AdamWOptimizer;
    updateOptimizer(config?: Partial<AdamWOptimizerConfig>): void;
    resumeFromLog(log: TrainingLogEntry): void;
    protected trainStep(state: Partial<TrainingState>, batch: {
        xs: Tensor;
        ys: Tensor;
    }, dummy?: boolean, keepGrads?: boolean): Scalar;
    private dummyPass;
    dispose(): void;
    private createEmptyState;
    stepDataset(dataset: Dataset<{
        xs: Tensor;
        ys: Tensor;
    }>, options: Partial<TrainingOptions>, validationDataset?: Dataset<{
        xs: Tensor;
        ys: Tensor;
    }>): Promise<{
        log: TrainingLogEntry;
    }>;
    private performLogging;
    trainOnDataset(dataset: Dataset<{
        xs: Tensor;
        ys: Tensor;
    }>, options: Partial<TrainingOptions>, validationDataset?: Dataset<{
        xs: Tensor;
        ys: Tensor;
    }>): Promise<{
        losses: number[];
        validationLosses: number[];
    }>;
}
