@trymirai/uzu 0.2.10 → 0.3.0

package/src/bridging/engine.ts

@@ -2,14 +2,16 @@ import { EngineError, EngineErrorCode, unwrapOrThrow } from '../error';
 import { EngineInteractor } from '../interactors/engineInteractor';
 import { Engine as NapiEngine } from '../napi/uzu';
 import { readEnv } from '../utilities/env';
-import { ChatModel } from './chatModel';
 import { ChatSession } from './chatSession';
+import { ClassificationSession } from './classificationSession';
 import { Config } from './config';
 import { DownloadHandle } from './downloadHandle';
 import { DownloadProgressUpdate } from './downloadProgressUpdate';
 import { DownloadPhase, DownloadState } from './downloadState';
 import { LicenseStatus, licenseStatusFromNapiLicenseStatus } from './licenseStatus';
-import { ModelType, modelTypeToNapiModelType } from './modelType';
+import { Model } from './model';
+import { ModelKind, modelKindToNapiModelKind } from './modelKind';
+import { TextToSpeechSession } from './textToSpeechSession';
 
 export class Engine {
     private readonly napiEngine: NapiEngine;
@@ -55,32 +57,18 @@ export class Engine {
         return new DownloadHandle(this.napiEngine.createModelDownloadHandle(repoId));
     }
 
-    async chatModels(types: ModelType[] | null = null): Promise<ChatModel[]> {
-        let typesToFilter: ModelType[] = types ?? [ModelType.Local, ModelType.Cloud];
-        let napiModelTypes = typesToFilter.map(modelTypeToNapiModelType);
-        let napiChatModels = await this.napiEngine.getChatModels(napiModelTypes);
-        let results: ChatModel[] = napiChatModels.map(ChatModel.fromNapiChatModel);
-        return results;
-    }
-
-    async chatModel(repoId: string, types: ModelType[] | null = null): Promise<ChatModel> {
-        const chatModels = await this.chatModels(types);
-        const chatModel = chatModels.find((model) => model.repoId === repoId);
-        return unwrapOrThrow(chatModel, EngineErrorCode.ModelNotFound);
-    }
-
-    async downloadChatModel(
-        chatModel: ChatModel,
+    async downloadModel(
+        model: Model,
         progress: (progressUpdate: DownloadProgressUpdate) => void = () => {},
     ): Promise<DownloadState> {
-        switch (chatModel.type) {
-            case ModelType.Local:
+        switch (model.kind) {
+            case ModelKind.Local:
                 break;
-            case ModelType.Cloud:
-                throw new EngineError(EngineErrorCode.UnexpectedModelType);
+            case ModelKind.Cloud:
+                throw new EngineError(EngineErrorCode.UnexpectedModelKind);
         }
 
-        const downloadHandle = this.downloadHandle(chatModel.repoId);
+        const downloadHandle = this.downloadHandle(model.repoId);
         const state = await downloadHandle.state();
         switch (state.phase) {
             case DownloadPhase.Downloaded:
@@ -97,9 +85,57 @@ export class Engine {
         return finalState;
     }
 
-    chatSession(chatModel: ChatModel, config: Config = Config.default()): ChatSession {
-        const napiChatModel = chatModel.toNapi();
+    async chatModels(kinds: ModelKind[] | null = null): Promise<Model[]> {
+        let kindsToFilter: ModelKind[] = kinds ?? [ModelKind.Local, ModelKind.Cloud];
+        let napiModelKinds = kindsToFilter.map(modelKindToNapiModelKind);
+        let napiChatModels = await this.napiEngine.getChatModels(napiModelKinds);
+        let results: Model[] = napiChatModels.map(Model.fromNapiModel);
+        return results;
+    }
+
+    async chatModel(repoId: string, kinds: ModelKind[] | null = null): Promise<Model> {
+        const chatModels = await this.chatModels(kinds);
+        const chatModel = chatModels.find((model) => model.repoId === repoId);
+        return unwrapOrThrow(chatModel, EngineErrorCode.ModelNotFound);
+    }
+
+    chatSession(model: Model, config: Config = Config.default()): ChatSession {
+        const napiModel = model.toNapi();
         const napiConfig = config.toNapi();
-        return new ChatSession(this.napiEngine.createChatSession(napiChatModel, napiConfig));
+        return new ChatSession(this.napiEngine.createChatSession(napiModel, napiConfig));
+    }
+
+    async classificationModels(): Promise<Model[]> {
+        const napiClassificationModels = await this.napiEngine.getClassificationModels();
+        const results: Model[] = napiClassificationModels.map(Model.fromNapiModel);
+        return results;
+    }
+
+    async classificationModel(repoId: string): Promise<Model> {
+        const classificationModels = await this.classificationModels();
+        const classificationModel = classificationModels.find((model) => model.repoId === repoId);
+        return unwrapOrThrow(classificationModel, EngineErrorCode.ModelNotFound);
+    }
+
+    classificationSession(model: Model): ClassificationSession {
+        const napiModel = model.toNapi();
+        return new ClassificationSession(this.napiEngine.createClassificationSession(napiModel));
+    }
+
+    async textToSpeechModels(): Promise<Model[]> {
+        const napiTextToSpeechModels = await this.napiEngine.getTextToSpeechModels();
+        const results: Model[] = napiTextToSpeechModels.map(Model.fromNapiModel);
+        return results;
+    }
+
+    async textToSpeechModel(repoId: string): Promise<Model> {
+        const textToSpeechModels = await this.textToSpeechModels();
+        const textToSpeechModel = textToSpeechModels.find((model) => model.repoId === repoId);
+        return unwrapOrThrow(textToSpeechModel, EngineErrorCode.ModelNotFound);
+    }
+
+    textToSpeechSession(model: Model): TextToSpeechSession {
+        const napiModel = model.toNapi();
+        return new TextToSpeechSession(this.napiEngine.createTextToSpeechSession(napiModel));
     }
 }

package/src/bridging/model.ts

@@ -0,0 +1,68 @@
+import { Model as NapiModel } from '../napi/uzu';
+import { ClassificationLabel } from './classificationLabel';
+import { ModelKind, modelKindFromNapiModelKind, modelKindToNapiModelKind } from './modelKind';
+import { ModelType, modelTypeFromNapiModelType, modelTypeToNapiModelType } from './modelType';
+
+export class Model {
+    readonly repoId: string;
+    readonly type: ModelType;
+    readonly kind: ModelKind;
+    readonly name: string;
+    readonly vendor: string;
+    readonly quantization: string | null;
+    readonly outputParserRegex: string | null;
+    readonly classificationLabels: ClassificationLabel[];
+
+    get isThinking(): boolean {
+        return this.outputParserRegex !== null;
+    }
+
+    private constructor(
+        repoId: string,
+        type: ModelType,
+        kind: ModelKind,
+        name: string,
+        vendor: string,
+        quantization: string | null,
+        outputParserRegex: string | null,
+        classificationLabels: ClassificationLabel[],
+    ) {
+        this.repoId = repoId;
+        this.type = type;
+        this.kind = kind;
+        this.name = name;
+        this.vendor = vendor;
+        this.quantization = quantization;
+        this.outputParserRegex = outputParserRegex;
+        this.classificationLabels = classificationLabels;
+    }
+
+    static fromNapiModel(napiModel: NapiModel): Model {
+        return new Model(
+            napiModel.repoId,
+            modelTypeFromNapiModelType(napiModel.type),
+            modelKindFromNapiModelKind(napiModel.kind),
+            napiModel.name,
+            napiModel.vendor,
+            napiModel.quantization ?? null,
+            napiModel.outputParserRegex ?? null,
+            napiModel.classificationLabels.map(ClassificationLabel.fromNapi),
+        );
+    }
+
+    toNapi(): NapiModel {
+        const napiModel: NapiModel = {
+            repoId: this.repoId,
+            type: modelTypeToNapiModelType(this.type),
+            kind: modelKindToNapiModelKind(this.kind),
+            name: this.name,
+            vendor: this.vendor,
+            ...(this.quantization !== null && { quantization: this.quantization }),
+            ...(this.outputParserRegex !== null && { outputParserRegex: this.outputParserRegex }),
+            classificationLabels: this.classificationLabels.map((classificationLabel) =>
+                classificationLabel.toNapi(),
+            ),
+        };
+        return napiModel;
+    }
+}

package/src/bridging/modelKind.ts

@@ -0,0 +1,27 @@
+import { EngineError, EngineErrorCode } from '../error';
+import { ModelKind as NapiModelKind } from '../napi/uzu';
+
+export enum ModelKind {
+    Local = 'Local',
+    Cloud = 'Cloud',
+}
+
+export function modelKindFromNapiModelKind(napiModelKind: NapiModelKind): ModelKind {
+    switch (napiModelKind) {
+        case 0:
+            return ModelKind.Local;
+        case 1:
+            return ModelKind.Cloud;
+        default:
+            throw new EngineError(EngineErrorCode.UnexpectedModelKind);
+    }
+}
+
+export function modelKindToNapiModelKind(modelKind: ModelKind): NapiModelKind {
+    switch (modelKind) {
+        case ModelKind.Local:
+            return 0;
+        case ModelKind.Cloud:
+            return 1;
+    }
+}

package/src/bridging/modelType.ts

@@ -2,16 +2,19 @@ import { EngineError, EngineErrorCode } from '../error';
 import { ModelType as NapiModelType } from '../napi/uzu';
 
 export enum ModelType {
-    Local = 'Local',
-    Cloud = 'Cloud',
+    Chat = 'Chat',
+    Classification = 'Classification',
+    TTS = 'TTS',
 }
 
 export function modelTypeFromNapiModelType(napiModelType: NapiModelType): ModelType {
     switch (napiModelType) {
         case 0:
-            return ModelType.Local;
+            return ModelType.Chat;
         case 1:
-            return ModelType.Cloud;
+            return ModelType.Classification;
+        case 2:
+            return ModelType.TTS;
         default:
             throw new EngineError(EngineErrorCode.UnexpectedModelType);
     }
@@ -19,9 +22,11 @@ export function modelTypeFromNapiModelType(napiModelType: NapiModelType): ModelT
 
 export function modelTypeToNapiModelType(modelType: ModelType): NapiModelType {
     switch (modelType) {
-        case ModelType.Local:
+        case ModelType.Chat:
             return 0;
-        case ModelType.Cloud:
+        case ModelType.Classification:
             return 1;
+        case ModelType.TTS:
+            return 2;
     }
 }

package/src/bridging/preset.ts

@@ -28,6 +28,20 @@ export class Preset implements ToNapi<NapiPreset> {
         return new Preset(napiPreset);
     }
 
+    static nGramSpeculator(useCase: string, numberOfSpeculatedTokens: number): Preset {
+        const napiPreset: NapiPreset = {
+            type: 'NGramSpeculator',
+            useCase,
+            numberOfSpeculatedTokens,
+        };
+        return new Preset(napiPreset);
+    }
+
+    static chat(): Preset {
+        const napiPreset: NapiPreset = { type: 'Chat' };
+        return new Preset(napiPreset);
+    }
+
     toNapi(): NapiPreset {
         return this.napiPreset;
     }

package/src/bridging/samplingMethod.ts

@@ -1,5 +1,9 @@
 import { SamplingMethod as NapiSamplingMethod } from '../napi/uzu';
 import { ToNapi } from './napi';
+import {
+    SamplingProcessingOrder,
+    samplingProcessingOrderToNapiSamplingProcessingOrder,
+} from './samplingProcessingOrder';
 
 export class SamplingMethod implements ToNapi<NapiSamplingMethod> {
     private readonly napiSamplingMethod: NapiSamplingMethod;
@@ -18,9 +22,11 @@ export class SamplingMethod implements ToNapi<NapiSamplingMethod> {
         topK: number | null,
         topP: number | null,
         minP: number | null,
+        processingOrder: SamplingProcessingOrder,
     ): SamplingMethod {
         let napiSamplingMethod: NapiSamplingMethod = {
             type: 'Stochastic',
+            processingOrder: samplingProcessingOrderToNapiSamplingProcessingOrder(processingOrder),
         };
         if (temperature !== null) {
             napiSamplingMethod.temperature = temperature;

package/src/bridging/samplingProcessingOrder.ts

@@ -0,0 +1,31 @@
+import { EngineError, EngineErrorCode } from '../error';
+import { SamplingProcessingOrder as NapiSamplingProcessingOrder } from '../napi/uzu';
+
+export enum SamplingProcessingOrder {
+    TemperatureThenFilters = 'TemperatureThenFilters',
+    FiltersThenTemperature = 'FiltersThenTemperature',
+}
+
+export function samplingProcessingOrderFromNapiSamplingProcessingOrder(
+    napiSamplingProcessingOrder: NapiSamplingProcessingOrder,
+): SamplingProcessingOrder {
+    switch (napiSamplingProcessingOrder) {
+        case 0:
+            return SamplingProcessingOrder.TemperatureThenFilters;
+        case 1:
+            return SamplingProcessingOrder.FiltersThenTemperature;
+        default:
+            throw new EngineError(EngineErrorCode.UnexpectedSamplingProcessingOrder);
+    }
+}
+
+export function samplingProcessingOrderToNapiSamplingProcessingOrder(
+    samplingProcessingOrder: SamplingProcessingOrder,
+): NapiSamplingProcessingOrder {
+    switch (samplingProcessingOrder) {
+        case SamplingProcessingOrder.TemperatureThenFilters:
+            return 0;
+        case SamplingProcessingOrder.FiltersThenTemperature:
+            return 1;
+    }
+}

package/src/bridging/textToSpeechSession.ts

@@ -0,0 +1,15 @@
+import { TextToSpeechSession as NapiTextToSpeechSession } from '../napi/uzu';
+import { Input } from './input';
+
+export class TextToSpeechSession {
+    private readonly napiTextToSpeechSession: NapiTextToSpeechSession;
+
+    constructor(napiTextToSpeechSession: NapiTextToSpeechSession) {
+        this.napiTextToSpeechSession = napiTextToSpeechSession;
+    }
+
+    run(input: Input, outputPath: string): void {
+        const napiInput = input.toNapi();
+        this.napiTextToSpeechSession.run(napiInput, outputPath);
+    }
+}

package/src/error.ts

@@ -6,6 +6,8 @@ export enum EngineErrorCode {
     UnexpectedFinishReason = 'unexpectedFinishReason',
     UnexpectedLicenseStatus = 'unexpectedLicenseStatus',
     UnexpectedModelType = 'unexpectedModelType',
+    UnexpectedModelKind = 'unexpectedModelKind',
+    UnexpectedSamplingProcessingOrder = 'unexpectedSamplingProcessingOrder',
     LicenseNotActivated = 'licenseNotActivated',
     Unknown = 'unknown',
 }
@@ -18,6 +20,8 @@ const errorCodeToMessage = {
     [EngineErrorCode.UnexpectedFinishReason]: 'Unexpected finish reason',
     [EngineErrorCode.UnexpectedLicenseStatus]: 'Unexpected license status',
     [EngineErrorCode.UnexpectedModelType]: 'Unexpected model type',
+    [EngineErrorCode.UnexpectedModelKind]: 'Unexpected model kind',
+    [EngineErrorCode.UnexpectedSamplingProcessingOrder]: 'Unexpected sampling processing order',
     [EngineErrorCode.LicenseNotActivated]: 'License not activated',
     [EngineErrorCode.Unknown]: 'Unknown error',
 };

package/src/index.ts

@@ -1,7 +1,10 @@
 export { AsyncBatchSize } from './bridging/asyncBatchSize';
-export { ChatModel } from './bridging/chatModel';
 export { ChatSession } from './bridging/chatSession';
 export { ClassificationFeature } from './bridging/classificationFeature';
+export { ClassificationLabel } from './bridging/classificationLabel';
+export { ClassificationOutput } from './bridging/classificationOutput';
+export { ClassificationSession } from './bridging/classificationSession';
+export { ClassificationStats } from './bridging/classificationStats';
 export { Config } from './bridging/config';
 export { ContextLength } from './bridging/contextLength';
 export { ContextMode } from './bridging/contextMode';
@@ -13,6 +16,8 @@ export { FinishReason } from './bridging/finishReason';
 export { GrammarConfig } from './bridging/grammarConfig';
 export { Input } from './bridging/input';
 export { Message } from './bridging/message';
+export { Model } from './bridging/model';
+export { ModelKind } from './bridging/modelKind';
 export { ModelType } from './bridging/modelType';
 export { Output } from './bridging/output';
 export { ParsedText } from './bridging/parsedText';
@@ -23,10 +28,12 @@ export { RunConfig } from './bridging/runConfig';
 export { RunStats } from './bridging/runStats';
 export { SamplingMethod } from './bridging/samplingMethod';
 export { SamplingPolicy } from './bridging/samplingPolicy';
+export { SamplingProcessingOrder } from './bridging/samplingProcessingOrder';
 export { SamplingSeed } from './bridging/samplingSeed';
 export { Stats } from './bridging/stats';
 export { StepStats } from './bridging/stepStats';
 export { Text } from './bridging/text';
+export { TextToSpeechSession } from './bridging/textToSpeechSession';
 export { TotalStats } from './bridging/totalStats';
 export { EngineError } from './error';
 export { ChatModelInteractor } from './interactors/chatModelInteractor';

package/src/interactors/chatModelInteractor.ts

@@ -1,12 +1,12 @@
 import { AsyncBatchSize } from '../bridging/asyncBatchSize';
-import { ChatModel } from '../bridging/chatModel';
 import { Config } from '../bridging/config';
 import { ContextLength } from '../bridging/contextLength';
 import { ContextMode } from '../bridging/contextMode';
 import { DownloadProgressUpdate } from '../bridging/downloadProgressUpdate';
 import { DownloadPhase } from '../bridging/downloadState';
 import { Message } from '../bridging/message';
-import { ModelType } from '../bridging/modelType';
+import { Model } from '../bridging/model';
+import { ModelKind } from '../bridging/modelKind';
 import { Output } from '../bridging/output';
 import { PrefillStepSize } from '../bridging/prefillStepSize';
 import { Preset } from '../bridging/preset';
@@ -17,18 +17,18 @@ import { ChatSessionInteractor } from './chatSessionInteractor';
 import { DownloadInteractor } from './downloadInteractor';
 import { Interactor, InteractorEntity } from './interactor';
 
-export class ChatModelInteractor implements Interactor<ChatModel> {
+export class ChatModelInteractor implements Interactor<Model> {
     readonly modelsInteractor: ChatModelsInteractor;
-    readonly entity: InteractorEntity<ChatModel>;
+    readonly entity: InteractorEntity<Model>;
     readonly config: Config;
 
-    constructor(modelsInteractor: ChatModelsInteractor, model: InteractorEntity<ChatModel>, config: Config) {
+    constructor(modelsInteractor: ChatModelsInteractor, model: InteractorEntity<Model>, config: Config) {
         this.modelsInteractor = modelsInteractor;
         this.entity = model;
         this.config = config;
     }
 
-    async finalize(): Promise<ChatModel> {
+    async finalize(): Promise<Model> {
         return await this.entity;
     }
 
@@ -103,7 +103,7 @@ export class ChatModelInteractor implements Interactor<ChatModel> {
             const engine = await this.modelsInteractor.engineInteractor.finalize();
 
             let model = await this.finalize();
-            if (model.type === ModelType.Local) {
+            if (model.kind === ModelKind.Local) {
                 model = await this.download().finalize();
             }
 

package/src/interactors/chatModelsInteractor.ts

@@ -1,20 +1,20 @@
-import { ChatModel } from '../bridging/chatModel';
 import { Config } from '../bridging/config';
+import { Model } from '../bridging/model';
 import { EngineErrorCode, unwrapOrThrow } from '../error';
 import { ChatModelInteractor } from './chatModelInteractor';
 import { EngineInteractor } from './engineInteractor';
 import { Interactor, InteractorEntity } from './interactor';
 
-export class ChatModelsInteractor implements Interactor<ChatModel[]> {
+export class ChatModelsInteractor implements Interactor<Model[]> {
     readonly engineInteractor: EngineInteractor;
-    readonly entity: InteractorEntity<ChatModel[]>;
+    readonly entity: InteractorEntity<Model[]>;
 
-    constructor(engineInteractor: EngineInteractor, models: InteractorEntity<ChatModel[]>) {
+    constructor(engineInteractor: EngineInteractor, models: InteractorEntity<Model[]>) {
         this.engineInteractor = engineInteractor;
         this.entity = models;
     }
 
-    async finalize(): Promise<ChatModel[]> {
+    async finalize(): Promise<Model[]> {
         return await this.entity;
     }
 

package/src/interactors/engineInteractor.ts

@@ -1,16 +1,16 @@
 import { Engine } from '../bridging/engine';
-import { ModelType } from '../bridging/modelType';
+import { ModelKind } from '../bridging/modelKind';
 import { ChatModelInteractor } from './chatModelInteractor';
 import { ChatModelsInteractor } from './chatModelsInteractor';
 import { Interactor, InteractorEntity } from './interactor';
 
 export class EngineInteractor implements Interactor<Engine> {
     readonly entity: InteractorEntity<Engine>;
-    readonly modelTypeFilter: ModelType[] | null;
+    readonly modelKindFilter: ModelKind[] | null;
 
-    constructor(engine: InteractorEntity<Engine>, modelTypeFilter: ModelType[] | null = null) {
+    constructor(engine: InteractorEntity<Engine>, modelKindFilter: ModelKind[] | null = null) {
         this.entity = engine;
-        this.modelTypeFilter = modelTypeFilter;
+        this.modelKindFilter = modelKindFilter;
     }
 
     async finalize(): Promise<Engine> {
@@ -19,14 +19,14 @@ export class EngineInteractor implements Interactor<Engine> {
 
     /* Models */
 
-    filterTypes(types: ModelType[] | null): EngineInteractor {
-        return new EngineInteractor(this.entity, types);
+    filterKinds(kinds: ModelKind[] | null): EngineInteractor {
+        return new EngineInteractor(this.entity, kinds);
     }
 
     chatModels(): ChatModelsInteractor {
         const modelsPromise = (async () => {
             const engine = await this.finalize();
-            return await engine.chatModels(this.modelTypeFilter);
+            return await engine.chatModels(this.modelKindFilter);
         })();
         return new ChatModelsInteractor(this, modelsPromise);
     }

package/src/napi/uzu.d.ts

@@ -5,17 +5,25 @@ export declare class ChatSession {
   reset(): void
 }
 
+export declare class ClassificationSession {
+  run(input: Input): ClassificationOutput
+}
+
 export declare class Engine {
   activate(apiKey: string): Promise<LicenseStatus>
-  getChatModels(types: Array<ModelType>): Promise<Array<ChatModel>>
+  getChatModels(kinds: Array<ModelKind>): Promise<Array<Model>>
+  getTextToSpeechModels(): Promise<Array<Model>>
+  getClassificationModels(): Promise<Array<Model>>
   getModelDownloadState(repoId: string): ModelDownloadState
   createModelDownloadHandle(repoId: string): ModelDownloadHandle
   benchmark(task: BenchmarksTask, prefillStepSize?: number | undefined | null): Promise<Array<BenchmarksResult>>
   constructor()
-  createChatSession(model: ChatModel, config: Config): ChatSession
+  createChatSession(model: Model, config: Config): ChatSession
+  createTextToSpeechSession(model: Model): TextToSpeechSession
+  createClassificationSession(model: Model): ClassificationSession
   registerLicenseStatusHandler(callback?: ((arg: LicenseStatus) => void) | undefined | null): void
   registerModelDownloadStateHandler(callback?: ((arg0: string, arg1: ModelDownloadState) => void) | undefined | null): void
-  registerChatModelsHandler(callback?: ((arg: Array<ChatModel>) => void) | undefined | null): void
+  registerModelsHandler(callback?: ((arg: Array<Model>) => void) | undefined | null): void
 }
 
 export declare class ModelDownloadHandle {
@@ -40,6 +48,10 @@ export declare class ProgressUpdate {
   progress: number
 }
 
+export declare class TextToSpeechSession {
+  run(input: Input, outputPath: string): void
+}
+
 export type AsyncBatchSize =
   | { type: 'Default' }
   | { type: 'Custom', size: number }
@@ -74,20 +86,36 @@ export interface BenchmarksTask {
   greedy: boolean
 }
 
-export interface ChatModel {
-  readonly repoId: string
-  readonly type: ModelType
-  readonly name: string
-  readonly vendor: string
-  readonly quantization?: string
-  readonly outputParserRegex?: string
-}
-
 export interface ClassificationFeature {
   name: string
   values: Array<string>
 }
 
+export interface ClassificationLabel {
+  index: number
+  label: string
+  title: string
+  description: string
+  recommendedThreshold: number
+}
+
+export interface ClassificationOutput {
+  logits: Array<number>
+  probabilities: Record<string, number>
+  stats: ClassificationStats
+}
+
+export interface ClassificationStats {
+  preprocessingDuration: number
+  forwardPassDuration: number
+  postprocessingDuration: number
+  totalDuration: number
+  tokensCount: number
+  tokensPerSecond: number
+  predictedLabel: string
+  confidence: number
+}
+
 export interface Config {
   preset: Preset
   contextMode: ContextMode
@@ -148,6 +176,17 @@ export interface Message {
   reasoningContent?: string
 }
 
+export interface Model {
+  readonly repoId: string
+  readonly type: ModelType
+  readonly kind: ModelKind
+  readonly name: string
+  readonly vendor: string
+  readonly quantization?: string
+  readonly outputParserRegex?: string
+  readonly classificationLabels: Array<ClassificationLabel>
+}
+
 export declare const enum ModelDownloadPhase {
   NotDownloaded = 0,
   Downloading = 1,
@@ -168,6 +207,11 @@ export interface ModelDownloadState {
   readonly error?: string
 }
 
+export declare const enum ModelKind {
+  Local = 0,
+  Cloud = 1
+}
+
 export type ModelStorageError =
   | { type: 'Io', message: string }
   | { type: 'Network', message: string }
@@ -179,8 +223,9 @@ export type ModelStorageError =
   | { type: 'Paused' }
 
 export declare const enum ModelType {
-  Local = 0,
-  Cloud = 1
+  Chat = 0,
+  Classification = 1,
+  TTS = 2
 }
 
 export interface Output {
@@ -203,6 +248,8 @@ export type Preset =
   | { type: 'General' }
   | { type: 'Classification', feature: ClassificationFeature }
   | { type: 'Summarization' }
+  | { type: 'NGramSpeculator', useCase: string, numberOfSpeculatedTokens: number }
+  | { type: 'Chat' }
 
 export declare const enum Role {
   System = 0,
@@ -224,12 +271,17 @@ export interface RunStats {
 
 export type SamplingMethod =
   | { type: 'Greedy' }
-  | { type: 'Stochastic', temperature?: number, topK?: number, topP?: number, minP?: number }
+  | { type: 'Stochastic', temperature?: number, topK?: number, topP?: number, minP?: number, processingOrder: SamplingProcessingOrder }
 
 export type SamplingPolicy =
   | { type: 'Default' }
   | { type: 'Custom', value: SamplingMethod }
 
+export declare const enum SamplingProcessingOrder {
+  TemperatureThenFilters = 0,
+  FiltersThenTemperature = 1
+}
+
 export type SamplingSeed =
   | { type: 'Default' }
   | { type: 'Custom', seed: number }
@@ -248,6 +300,8 @@ export interface StepStats {
   processedTokensPerSecond: number
   modelRun: RunStats
   run?: RunStats
+  speculatorProposed: bigint
+  speculatorAccepted: bigint
 }
 
 export type StorageError =