Skip to content

Add LiteRT trait and LiteRTLanguageModel backend#175

Open
john-rocky wants to merge 1 commit into
huggingface:mainfrom
john-rocky:litert-backend
Open

Add LiteRT trait and LiteRTLanguageModel backend#175
john-rocky wants to merge 1 commit into
huggingface:mainfrom
john-rocky:litert-backend

Conversation

@john-rocky

Copy link
Copy Markdown

Implements #174: a LiteRT package trait and a LiteRTLanguageModel backend that runs .litertlm models (Gemma 4 E2B/E4B, or any .litertlm on the Hub) fully on-device via Google's LiteRT-LM runtime, with Metal GPU acceleration on iOS and macOS.

let model = LiteRTLanguageModel(model: .gemma4_E2B)
let session = LanguageModelSession(model: model)
let response = try await session.respond(to: "What is the capital of France?")

What's included

  • Package.swift: a LiteRT trait; the john-rocky/swift-litert-lm dependency is gated by .when(platforms: [.iOS, .macOS], traits: ["LiteRT"]), so default builds and other platforms are unaffected.
  • LiteRTLanguageModel (behind #if LiteRT), following the existing backend patterns:
    • Three initializers: catalog model (auto-download from the Hub), any Hugging Face repo hosting a .litertlm, or a local file. Engine bring-up is lazy (first request or prewarm), matching how MLXLanguageModel defers loading.
    • The session transcript maps to a LiteRT-LM conversation per request (system message from .instructions + prior turns as initialMessages + trigger message), so multi-turn sessions and multiple sessions per model instance behave correctly.
    • respond / streamResponse for text; image segments (data, file URL, or fetched remote URL) for models with a vision tower.
    • Structured generation via schema-in-prompt + balanced-JSON extraction (honors includeSchemaInPrompt).
    • Tool calling for respond: prompt-driven tool descriptions, tool-call parsing, and the standard ToolExecutionDecision delegate flow, with the result fed back for a final answer (capped at 4 rounds).
    • GenerationOptions mapping: temperature, .greedy, .topK, .nucleus → LiteRT-LM's sampler config.
  • Tests: env-gated suite (LITERT_TEST_MODEL=/path/to/model.litertlm), same spirit as the MLX suite — the weights are multi-GB, so they don't run in CI.
  • README: provider list, trait docs, image-support table, and a LiteRT-LM usage section.

Verification

  • swift build --traits LiteRT — builds clean (no warnings) on macOS arm64.
  • swift build (default traits) — unaffected; the new file is fully #if LiteRT-gated and the dependency isn't resolved into the build graph without the trait.
  • swift format lint — clean on the new files.

Known limitations (documented in the README section)

Runs .litertlm models (e.g. Gemma 4) fully on-device via Google's
LiteRT-LM runtime, with Metal GPU acceleration on iOS and macOS.

- LiteRT package trait, gating a swift-litert-lm dependency to
  iOS/macOS; default builds are unaffected
- LiteRTLanguageModel: respond/streamResponse, image inputs for
  models with a vision tower, structured generation via
  schema-in-prompt + JSON extraction, and prompt-driven tool calling
  for respond (with the ToolExecutionDecision delegate flow)
- Env-gated tests (LITERT_TEST_MODEL), README section and provider
  table updates
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant