import { Tensor } from '@tensorflow/tfjs-core';
import { Conversation, ITokeniser } from '../tokeniser/type';
import { Dataset } from '@tensorflow/tfjs-data';
export declare const PAGE_FACTOR = 8;
export declare function flattenTokens(textData: Conversation[][], tokenizer: ITokeniser): Promise<Uint16Array>;
export declare function shuffle(array: Uint32Array): Uint32Array;
export interface DatasetState {
    shuffledIndexes: Uint32Array;
    step: number;
}
export declare class DatasetBuilder {
    tokenizer: ITokeniser;
    blockSize: number;
    constructor(tokenizer: ITokeniser, blockSize?: number);
    createTextDataset(flatTokens: Uint16Array, batchSize?: number, indexes?: Uint32Array): Promise<{
        dataset: Dataset<{
            xs: Tensor;
            ys: Tensor;
        }>;
        state: DatasetState;
    }>;
}
