@inferrlm/react-native-mlx 0.2.0-inferrlm.2 → 0.4.0

package/src/models.ts

@@ -1,3 +1,39 @@
+export enum ModelFamily {
+  Llama = 'Llama',
+  Qwen = 'Qwen',
+  Gemma = 'Gemma',
+  Phi = 'Phi',
+  SmolLM = 'SmolLM',
+  OpenELM = 'OpenELM',
+  PocketTTS = 'PocketTTS',
+  GLMASR = 'GLMASR',
+}
+
+export enum ModelProvider {
+  Meta = 'Meta',
+  Alibaba = 'Alibaba',
+  Google = 'Google',
+  Microsoft = 'Microsoft',
+  HuggingFace = 'HuggingFace',
+  Apple = 'Apple',
+  Kyutai = 'Kyutai',
+}
+
+export type ModelQuantization = '4bit' | '8bit' | 'bf16'
+
+export type ModelType = 'llm' | 'tts' | 'stt'
+
+export interface ModelInfo {
+  id: MLXModel
+  family: ModelFamily
+  provider: ModelProvider
+  parameters: string
+  quantization: ModelQuantization
+  displayName: string
+  downloadSize: number
+  type: ModelType
+}
+
 export enum MLXModel {
   // Llama 3.2 (Meta) - 1B and 3B variants
   Llama_3_2_1B_Instruct_4bit = 'mlx-community/Llama-3.2-1B-Instruct-4bit',
@@ -42,4 +78,315 @@ export enum MLXModel {
   OpenELM_1_1B_8bit = 'mlx-community/OpenELM-1_1B-8bit',
   OpenELM_3B_4bit = 'mlx-community/OpenELM-3B-4bit',
   OpenELM_3B_8bit = 'mlx-community/OpenELM-3B-8bit',
+
+  // PocketTTS (Kyutai) - 44.6M
+  PocketTTS = 'mlx-community/pocket-tts',
+  PocketTTS_8bit = 'mlx-community/pocket-tts-8bit',
+  PocketTTS_4bit = 'mlx-community/pocket-tts-4bit',
+
+  // GLM-ASR (GLMASR) - Speech-to-Text
+  GLM_ASR_Nano_4bit = 'mlx-community/GLM-ASR-Nano-2512-4bit',
 }
+
+export const MLXModels: ModelInfo[] = [
+  {
+    id: MLXModel.Llama_3_2_1B_Instruct_4bit,
+    family: ModelFamily.Llama,
+    provider: ModelProvider.Meta,
+    parameters: '1B',
+    quantization: '4bit',
+    displayName: 'Llama 3.2 1B Instruct (4-bit)',
+    downloadSize: 1407777762,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Llama_3_2_1B_Instruct_8bit,
+    family: ModelFamily.Llama,
+    provider: ModelProvider.Meta,
+    parameters: '1B',
+    quantization: '8bit',
+    displayName: 'Llama 3.2 1B Instruct (8-bit)',
+    downloadSize: 1313157436,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Llama_3_2_3B_Instruct_4bit,
+    family: ModelFamily.Llama,
+    provider: ModelProvider.Meta,
+    parameters: '3B',
+    quantization: '4bit',
+    displayName: 'Llama 3.2 3B Instruct (4-bit)',
+    downloadSize: 2019397474,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Llama_3_2_3B_Instruct_8bit,
+    family: ModelFamily.Llama,
+    provider: ModelProvider.Meta,
+    parameters: '3B',
+    quantization: '8bit',
+    displayName: 'Llama 3.2 3B Instruct (8-bit)',
+    downloadSize: 3413784042,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Qwen2_5_0_5B_Instruct_4bit,
+    family: ModelFamily.Qwen,
+    provider: ModelProvider.Alibaba,
+    parameters: '0.5B',
+    quantization: '4bit',
+    displayName: 'Qwen 2.5 0.5B Instruct (4-bit)',
+    downloadSize: 278064920,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Qwen2_5_0_5B_Instruct_8bit,
+    family: ModelFamily.Qwen,
+    provider: ModelProvider.Alibaba,
+    parameters: '0.5B',
+    quantization: '8bit',
+    displayName: 'Qwen 2.5 0.5B Instruct (8-bit)',
+    downloadSize: 525045902,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Qwen2_5_1_5B_Instruct_4bit,
+    family: ModelFamily.Qwen,
+    provider: ModelProvider.Alibaba,
+    parameters: '1.5B',
+    quantization: '4bit',
+    displayName: 'Qwen 2.5 1.5B Instruct (4-bit)',
+    downloadSize: 868628559,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Qwen2_5_1_5B_Instruct_8bit,
+    family: ModelFamily.Qwen,
+    provider: ModelProvider.Alibaba,
+    parameters: '1.5B',
+    quantization: '8bit',
+    displayName: 'Qwen 2.5 1.5B Instruct (8-bit)',
+    downloadSize: 1640414038,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Qwen2_5_3B_Instruct_4bit,
+    family: ModelFamily.Qwen,
+    provider: ModelProvider.Alibaba,
+    parameters: '3B',
+    quantization: '4bit',
+    displayName: 'Qwen 2.5 3B Instruct (4-bit)',
+    downloadSize: 1736293090,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Qwen2_5_3B_Instruct_8bit,
+    family: ModelFamily.Qwen,
+    provider: ModelProvider.Alibaba,
+    parameters: '3B',
+    quantization: '8bit',
+    displayName: 'Qwen 2.5 3B Instruct (8-bit)',
+    downloadSize: 3279142142,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Qwen3_1_7B_4bit,
+    family: ModelFamily.Qwen,
+    provider: ModelProvider.Alibaba,
+    parameters: '1.7B',
+    quantization: '4bit',
+    displayName: 'Qwen 3 1.7B (4-bit)',
+    downloadSize: 979502864,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Qwen3_1_7B_8bit,
+    family: ModelFamily.Qwen,
+    provider: ModelProvider.Alibaba,
+    parameters: '1.7B',
+    quantization: '8bit',
+    displayName: 'Qwen 3 1.7B (8-bit)',
+    downloadSize: 1839729195,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Gemma_3_1B_IT_4bit,
+    family: ModelFamily.Gemma,
+    provider: ModelProvider.Google,
+    parameters: '1B',
+    quantization: '4bit',
+    displayName: 'Gemma 3 1B IT (4-bit)',
+    downloadSize: 770650946,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Gemma_3_1B_IT_8bit,
+    family: ModelFamily.Gemma,
+    provider: ModelProvider.Google,
+    parameters: '1B',
+    quantization: '8bit',
+    displayName: 'Gemma 3 1B IT (8-bit)',
+    downloadSize: 1421522471,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Phi_3_5_Mini_Instruct_4bit,
+    family: ModelFamily.Phi,
+    provider: ModelProvider.Microsoft,
+    parameters: '3.8B',
+    quantization: '4bit',
+    displayName: 'Phi 3.5 Mini Instruct (4-bit)',
+    downloadSize: 2150195856,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Phi_3_5_Mini_Instruct_8bit,
+    family: ModelFamily.Phi,
+    provider: ModelProvider.Microsoft,
+    parameters: '3.8B',
+    quantization: '8bit',
+    displayName: 'Phi 3.5 Mini Instruct (8-bit)',
+    downloadSize: 4060636056,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Phi_4_Mini_Instruct_4bit,
+    family: ModelFamily.Phi,
+    provider: ModelProvider.Microsoft,
+    parameters: '3.8B',
+    quantization: '4bit',
+    displayName: 'Phi 4 Mini Instruct (4-bit)',
+    downloadSize: 2173624891,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.Phi_4_Mini_Instruct_8bit,
+    family: ModelFamily.Phi,
+    provider: ModelProvider.Microsoft,
+    parameters: '3.8B',
+    quantization: '8bit',
+    displayName: 'Phi 4 Mini Instruct (8-bit)',
+    downloadSize: 4091536167,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.SmolLM_1_7B_Instruct_4bit,
+    family: ModelFamily.SmolLM,
+    provider: ModelProvider.HuggingFace,
+    parameters: '1.7B',
+    quantization: '4bit',
+    displayName: 'SmolLM 1.7B Instruct (4-bit)',
+    downloadSize: 962855374,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.SmolLM_1_7B_Instruct_8bit,
+    family: ModelFamily.SmolLM,
+    provider: ModelProvider.HuggingFace,
+    parameters: '1.7B',
+    quantization: '8bit',
+    displayName: 'SmolLM 1.7B Instruct (8-bit)',
+    downloadSize: 1818493993,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.SmolLM2_1_7B_Instruct_4bit,
+    family: ModelFamily.SmolLM,
+    provider: ModelProvider.HuggingFace,
+    parameters: '1.7B',
+    quantization: '4bit',
+    displayName: 'SmolLM2 1.7B Instruct (4-bit)',
+    downloadSize: 980000000,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.SmolLM2_1_7B_Instruct_8bit,
+    family: ModelFamily.SmolLM,
+    provider: ModelProvider.HuggingFace,
+    parameters: '1.7B',
+    quantization: '8bit',
+    displayName: 'SmolLM2 1.7B Instruct (8-bit)',
+    downloadSize: 1850000000,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.OpenELM_1_1B_4bit,
+    family: ModelFamily.OpenELM,
+    provider: ModelProvider.Apple,
+    parameters: '1.1B',
+    quantization: '4bit',
+    displayName: 'OpenELM 1.1B (4-bit)',
+    downloadSize: 608162655,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.OpenELM_1_1B_8bit,
+    family: ModelFamily.OpenELM,
+    provider: ModelProvider.Apple,
+    parameters: '1.1B',
+    quantization: '8bit',
+    displayName: 'OpenELM 1.1B (8-bit)',
+    downloadSize: 1148048397,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.OpenELM_3B_4bit,
+    family: ModelFamily.OpenELM,
+    provider: ModelProvider.Apple,
+    parameters: '3B',
+    quantization: '4bit',
+    displayName: 'OpenELM 3B (4-bit)',
+    downloadSize: 1650000000,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.OpenELM_3B_8bit,
+    family: ModelFamily.OpenELM,
+    provider: ModelProvider.Apple,
+    parameters: '3B',
+    quantization: '8bit',
+    displayName: 'OpenELM 3B (8-bit)',
+    downloadSize: 3100000000,
+    type: 'llm',
+  },
+  {
+    id: MLXModel.PocketTTS,
+    family: ModelFamily.PocketTTS,
+    provider: ModelProvider.Kyutai,
+    parameters: '44.6M',
+    quantization: 'bf16',
+    displayName: 'PocketTTS (bf16)',
+    downloadSize: 180000000,
+    type: 'tts',
+  },
+  {
+    id: MLXModel.PocketTTS_8bit,
+    family: ModelFamily.PocketTTS,
+    provider: ModelProvider.Kyutai,
+    parameters: '44.6M',
+    quantization: '8bit',
+    displayName: 'PocketTTS (8-bit)',
+    downloadSize: 140000000,
+    type: 'tts',
+  },
+  {
+    id: MLXModel.PocketTTS_4bit,
+    family: ModelFamily.PocketTTS,
+    provider: ModelProvider.Kyutai,
+    parameters: '44.6M',
+    quantization: '4bit',
+    displayName: 'PocketTTS (4-bit)',
+    downloadSize: 80000000,
+    type: 'tts',
+  },
+  {
+    id: MLXModel.GLM_ASR_Nano_4bit,
+    family: ModelFamily.GLMASR,
+    provider: ModelProvider.Alibaba,
+    parameters: '1B',
+    quantization: '4bit',
+    displayName: 'GLM-ASR Nano (4-bit)',
+    downloadSize: 600000000,
+    type: 'stt',
+  },
+]

package/src/specs/LLM.nitro.ts

@@ -1,24 +1,109 @@
-import type { HybridObject } from 'react-native-nitro-modules'
+import type { AnyMap, HybridObject } from 'react-native-nitro-modules'
 
 /**
  * Statistics from the last text generation.
  */
 export interface GenerationStats {
-  /** Total number of tokens generated */
   tokenCount: number
-  /** Generation speed in tokens per second */
   tokensPerSecond: number
-  /** Time in milliseconds until the first token was generated */
   timeToFirstToken: number
-  /** Total generation time in milliseconds */
   totalTime: number
+  toolExecutionTime: number
 }
 
+export interface GenerationStartEvent {
+  type: 'generation_start'
+  timestamp: number
+}
+
+export interface TokenEvent {
+  type: 'token'
+  token: string
+}
+
+export interface ThinkingStartEvent {
+  type: 'thinking_start'
+  timestamp: number
+}
+
+export interface ThinkingChunkEvent {
+  type: 'thinking_chunk'
+  chunk: string
+}
+
+export interface ThinkingEndEvent {
+  type: 'thinking_end'
+  content: string
+  timestamp: number
+}
+
+export interface ToolCallStartEvent {
+  type: 'tool_call_start'
+  id: string
+  name: string
+  arguments: string
+}
+
+export interface ToolCallExecutingEvent {
+  type: 'tool_call_executing'
+  id: string
+}
+
+export interface ToolCallCompletedEvent {
+  type: 'tool_call_completed'
+  id: string
+  result: string
+}
+
+export interface ToolCallFailedEvent {
+  type: 'tool_call_failed'
+  id: string
+  error: string
+}
+
+export interface GenerationEndEvent {
+  type: 'generation_end'
+  content: string
+  stats: GenerationStats
+}
+
+export type StreamEvent =
+  | GenerationStartEvent
+  | TokenEvent
+  | ThinkingStartEvent
+  | ThinkingChunkEvent
+  | ThinkingEndEvent
+  | ToolCallStartEvent
+  | ToolCallExecutingEvent
+  | ToolCallCompletedEvent
+  | ToolCallFailedEvent
+  | GenerationEndEvent
+
 export interface LLMMessage {
   role: string
   content: string
 }
 
+/**
+ * Parameter definition for a tool.
+ */
+export interface ToolParameter {
+  name: string
+  type: string
+  description: string
+  required: boolean
+}
+
+/**
+ * Tool definition that can be called by the model.
+ */
+export interface ToolDefinition {
+  name: string
+  description: string
+  parameters: ToolParameter[]
+  handler: (args: AnyMap) => Promise<AnyMap>
+}
+
 /** Options for loading a model.
  */
 export interface LLMLoadOptions {
@@ -28,6 +113,8 @@ export interface LLMLoadOptions {
   additionalContext?: LLMMessage[]
   /** Whether to automatically manage message history */
   manageHistory?: boolean
+  /** Tools available for the model to call */
+  tools?: ToolDefinition[]
 }
 
 /**
@@ -50,12 +137,23 @@ export interface LLM extends HybridObject<{ ios: 'swift' }> {
   generate(prompt: string): Promise<string>
 
   /**
-   * Stream a response token by token.
+   * Stream a response token by token with optional tool calling support.
+   * Tools are automatically executed when the model calls them.
    * @param prompt - The input text to generate a response for
    * @param onToken - Callback invoked for each generated token
+   * @param onToolCall - Optional callback invoked when a tool is called (for UI feedback)
    * @returns The complete generated text
    */
-  stream(prompt: string, onToken: (token: string) => void): Promise<string>
+  stream(
+    prompt: string,
+    onToken: (token: string) => void,
+    onToolCall?: (toolName: string, args: string) => void,
+  ): Promise<string>
+
+  streamWithEvents(
+    prompt: string,
+    onEvent: (eventJson: string) => void,
+  ): Promise<string>
 
   /**
    * Stop the current generation.
@@ -96,3 +194,9 @@ export interface LLM extends HybridObject<{ ios: 'swift' }> {
   /** System prompt used when loading the model */
   systemPrompt: string
 }
+
+/**
+ * Supported parameter types for tool definitions.
+ * Used for type safety in createTool().
+ */
+export type ToolParameterType = 'string' | 'number' | 'boolean' | 'array' | 'object'

package/src/specs/STT.nitro.ts

@@ -0,0 +1,35 @@
+import type { HybridObject } from 'react-native-nitro-modules'
+
+export interface STTLoadOptions {
+  onProgress?: (progress: number) => void
+}
+
+export interface STTTranscriptionInfo {
+  promptTokens: number
+  generationTokens: number
+  tokensPerSecond: number
+  prefillTime: number
+  generateTime: number
+}
+
+export interface STT extends HybridObject<{ ios: 'swift' }> {
+  readonly isLoaded: boolean
+  readonly isTranscribing: boolean
+  readonly isListening: boolean
+  readonly modelId: string
+
+  load(modelId: string, options?: STTLoadOptions): Promise<void>
+
+  transcribe(audio: ArrayBuffer): Promise<string>
+  transcribeStream(
+    audio: ArrayBuffer,
+    onToken: (token: string) => void
+  ): Promise<string>
+
+  startListening(): Promise<void>
+  transcribeBuffer(): Promise<string>
+  stopListening(): Promise<string>
+
+  stop(): void
+  unload(): void
+}

package/src/specs/TTS.nitro.ts

@@ -0,0 +1,30 @@
+import type { HybridObject } from 'react-native-nitro-modules'
+
+export interface TTSLoadOptions {
+  onProgress?: (progress: number) => void
+}
+
+export interface TTSGenerateOptions {
+  voice?: string
+  speed?: number
+}
+
+export interface TTS extends HybridObject<{ ios: 'swift' }> {
+  readonly isLoaded: boolean
+  readonly isGenerating: boolean
+  readonly modelId: string
+  readonly sampleRate: number
+
+  load(modelId: string, options?: TTSLoadOptions): Promise<void>
+  generate(
+    text: string,
+    options?: TTSGenerateOptions
+  ): Promise<ArrayBuffer>
+  stream(
+    text: string,
+    onAudioChunk: (audio: ArrayBuffer) => void,
+    options?: TTSGenerateOptions
+  ): Promise<void>
+  stop(): void
+  unload(): void
+}

package/src/stt.ts

@@ -0,0 +1,67 @@
+import { NitroModules } from 'react-native-nitro-modules'
+import type {
+  STT as STTSpec,
+  STTLoadOptions,
+} from './specs/STT.nitro'
+
+let instance: STTSpec | null = null
+
+function getInstance(): STTSpec {
+  if (!instance) {
+    instance = NitroModules.createHybridObject<STTSpec>('STT')
+  }
+  return instance
+}
+
+export const STT = {
+  load(modelId: string, options?: STTLoadOptions): Promise<void> {
+    return getInstance().load(modelId, options)
+  },
+
+  transcribe(audio: ArrayBuffer): Promise<string> {
+    return getInstance().transcribe(audio)
+  },
+
+  transcribeStream(
+    audio: ArrayBuffer,
+    onToken: (token: string) => void
+  ): Promise<string> {
+    return getInstance().transcribeStream(audio, onToken)
+  },
+
+  startListening(): Promise<void> {
+    return getInstance().startListening()
+  },
+
+  transcribeBuffer(): Promise<string> {
+    return getInstance().transcribeBuffer()
+  },
+
+  stopListening(): Promise<string> {
+    return getInstance().stopListening()
+  },
+
+  stop(): void {
+    getInstance().stop()
+  },
+
+  unload(): void {
+    getInstance().unload()
+  },
+
+  get isLoaded(): boolean {
+    return getInstance().isLoaded
+  },
+
+  get isTranscribing(): boolean {
+    return getInstance().isTranscribing
+  },
+
+  get isListening(): boolean {
+    return getInstance().isListening
+  },
+
+  get modelId(): string {
+    return getInstance().modelId
+  },
+}