feat: add managed Granite embedding models
GoodBuddy previously required an external OpenAI-compatible endpoint for vector retrieval. It can now install, verify, import, remove, and run the built-in Granite multilingual embedding model in an isolated local inference process, while retaining custom embedding connections and migrating existing settings. The model connection UI now manages built-in and custom vector connections consistently, reports the actual retrieval path, preserves credentials across failed saves, and renders the vector tab synchronously. Related settings copy, switch styling, metadata typography, and renderer bundle guardrails are aligned with the resulting workflow. Release note: 新增可管理的内置 Granite 多语言向量模型，支持本地下载、离线导入、连接切换与诊断；同时优化向量设置及相关界面的文案、开关和状态显示。
This commit is contained in:
mesalogo
@@ -0,0 +1,400 @@
|
||||
import { readFile } from 'node:fs/promises'
|
||||
import path from 'node:path'
|
||||
import { Tokenizer } from '@huggingface/tokenizers'
|
||||
import * as ort from 'onnxruntime-web/wasm'
|
||||
import type {
|
||||
EmbeddingInferenceEngine,
|
||||
EmbeddingOnnxRuntime,
|
||||
EmbeddingTokenizer
|
||||
} from './embedding-inference-worker'
|
||||
|
||||
const MODEL_FILE = 'model_quantized.onnx'
|
||||
const TOKENIZER_FILE = 'tokenizer.json'
|
||||
const TOKENIZER_CONFIG_FILE = 'tokenizer_config.json'
|
||||
const MAX_TOKENS = 32_768
|
||||
const EMBEDDING_DIMENSIONS = 384
|
||||
|
||||
type Encoding = {
|
||||
readonly ids: readonly number[]
|
||||
readonly attention_mask: readonly number[]
|
||||
readonly token_type_ids?: readonly number[]
|
||||
}
|
||||
|
||||
export interface GraniteTokenizerLike {
|
||||
encode(text: string, options?: {
|
||||
readonly add_special_tokens?: boolean
|
||||
readonly return_token_type_ids?: boolean | null
|
||||
}): Encoding
|
||||
token_to_id(token: string): number | undefined
|
||||
}
|
||||
|
||||
export interface GraniteTensorLike {
|
||||
readonly type: string
|
||||
readonly dims: readonly number[]
|
||||
readonly data: ArrayLike<number>
|
||||
}
|
||||
|
||||
export interface GraniteSessionLike {
|
||||
readonly inputNames: readonly string[]
|
||||
readonly outputNames: readonly string[]
|
||||
run(
|
||||
feeds: Readonly<Record<string, unknown>>
|
||||
): Promise<Readonly<Record<string, GraniteTensorLike>>>
|
||||
release?(): void
|
||||
}
|
||||
|
||||
type LoadedEngine = {
|
||||
readonly tokenizer: GraniteTokenizerLike
|
||||
readonly padTokenId: number
|
||||
readonly session: GraniteSessionLike
|
||||
}
|
||||
|
||||
export interface GraniteEmbeddingEngineDependencies {
|
||||
readonly loadTokenizer?: (
|
||||
tokenizerJson: unknown,
|
||||
tokenizerConfig: unknown
|
||||
) => GraniteTokenizerLike
|
||||
readonly createSession?: (model: Uint8Array) => Promise<GraniteSessionLike>
|
||||
readonly createInt64Tensor?: (
|
||||
data: BigInt64Array,
|
||||
dimensions: readonly [number, number]
|
||||
) => unknown
|
||||
readonly readFile?: (filePath: string) => Promise<Uint8Array>
|
||||
}
|
||||
|
||||
type TokenizedBatch = {
|
||||
readonly batchSize: number
|
||||
readonly sequenceLength: number
|
||||
readonly inputIds: BigInt64Array
|
||||
readonly attentionMask: BigInt64Array
|
||||
readonly tokenTypeIds?: BigInt64Array
|
||||
}
|
||||
|
||||
function abortIfRequested(signal: AbortSignal): void {
|
||||
if (signal.aborted) {
|
||||
throw signal.reason ?? new Error('Embedding inference was cancelled')
|
||||
}
|
||||
}
|
||||
|
||||
function asJsonObject(value: unknown, name: string): Record<string, unknown> {
|
||||
if (typeof value !== 'object' || value === null || Array.isArray(value)) {
|
||||
throw new TypeError(`${name} must contain a JSON object`)
|
||||
}
|
||||
return value as Record<string, unknown>
|
||||
}
|
||||
|
||||
function resolvePadTokenId(
|
||||
tokenizer: GraniteTokenizerLike,
|
||||
rawConfig: unknown
|
||||
): number {
|
||||
const config = asJsonObject(rawConfig, TOKENIZER_CONFIG_FILE)
|
||||
const configuredId = config.pad_token_id
|
||||
if (
|
||||
typeof configuredId === 'number' &&
|
||||
Number.isSafeInteger(configuredId) &&
|
||||
configuredId >= 0
|
||||
) {
|
||||
return configuredId
|
||||
}
|
||||
if (typeof config.pad_token === 'string') {
|
||||
const tokenId = tokenizer.token_to_id(config.pad_token)
|
||||
if (tokenId !== undefined) {
|
||||
return tokenId
|
||||
}
|
||||
}
|
||||
throw new Error('Tokenizer configuration does not define a valid pad token')
|
||||
}
|
||||
|
||||
function parseJson(bytes: Uint8Array, name: string): unknown {
|
||||
try {
|
||||
return JSON.parse(new TextDecoder().decode(bytes)) as unknown
|
||||
} catch {
|
||||
throw new Error(`${name} is not valid JSON`)
|
||||
}
|
||||
}
|
||||
|
||||
export function buildGraniteInputBatch(options: {
|
||||
readonly tokenizer: GraniteTokenizerLike
|
||||
readonly texts: readonly string[]
|
||||
readonly padTokenId: number
|
||||
readonly includeTokenTypeIds: boolean
|
||||
}): TokenizedBatch {
|
||||
if (options.texts.length === 0) {
|
||||
throw new RangeError('At least one text is required')
|
||||
}
|
||||
const encodings = options.texts.map((text) =>
|
||||
options.tokenizer.encode(text, {
|
||||
add_special_tokens: true,
|
||||
return_token_type_ids: options.includeTokenTypeIds
|
||||
})
|
||||
)
|
||||
const sequenceLength = Math.max(...encodings.map(({ ids }) => ids.length))
|
||||
if (sequenceLength > MAX_TOKENS) {
|
||||
throw new RangeError(
|
||||
`Embedding input exceeds the ${MAX_TOKENS} token limit`
|
||||
)
|
||||
}
|
||||
if (sequenceLength < 1) {
|
||||
throw new RangeError('Tokenizer produced an empty input')
|
||||
}
|
||||
|
||||
const itemCount = options.texts.length * sequenceLength
|
||||
const inputIds = new BigInt64Array(itemCount)
|
||||
inputIds.fill(BigInt(options.padTokenId))
|
||||
const attentionMask = new BigInt64Array(itemCount)
|
||||
const tokenTypeIds = options.includeTokenTypeIds
|
||||
? new BigInt64Array(itemCount)
|
||||
: undefined
|
||||
|
||||
for (const [batchIndex, encoding] of encodings.entries()) {
|
||||
if (
|
||||
encoding.attention_mask.length !== encoding.ids.length ||
|
||||
(options.includeTokenTypeIds &&
|
||||
encoding.token_type_ids !== undefined &&
|
||||
encoding.token_type_ids.length !== encoding.ids.length)
|
||||
) {
|
||||
throw new Error('Tokenizer produced inconsistent input lengths')
|
||||
}
|
||||
const offset = batchIndex * sequenceLength
|
||||
for (let index = 0; index < encoding.ids.length; index += 1) {
|
||||
const tokenId = encoding.ids[index]
|
||||
const mask = encoding.attention_mask[index]
|
||||
if (
|
||||
tokenId === undefined ||
|
||||
mask === undefined ||
|
||||
!Number.isSafeInteger(tokenId) ||
|
||||
!Number.isSafeInteger(mask)
|
||||
) {
|
||||
throw new Error('Tokenizer produced an invalid token')
|
||||
}
|
||||
inputIds[offset + index] = BigInt(tokenId)
|
||||
attentionMask[offset + index] = BigInt(mask)
|
||||
if (tokenTypeIds) {
|
||||
tokenTypeIds[offset + index] = BigInt(
|
||||
encoding.token_type_ids?.[index] ?? 0
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
batchSize: options.texts.length,
|
||||
sequenceLength,
|
||||
inputIds,
|
||||
attentionMask,
|
||||
tokenTypeIds
|
||||
}
|
||||
}
|
||||
|
||||
export function poolGraniteClsEmbeddings(options: {
|
||||
readonly outputs: Readonly<Record<string, GraniteTensorLike>>
|
||||
readonly outputNames: readonly string[]
|
||||
readonly batchSize: number
|
||||
readonly sequenceLength: number
|
||||
}): number[][] {
|
||||
const preferred = options.outputs.last_hidden_state
|
||||
const output =
|
||||
preferred?.dims.length === 3 &&
|
||||
(preferred.type === 'float32' || preferred.type === 'float64')
|
||||
? preferred
|
||||
: options.outputNames
|
||||
.map((name) => options.outputs[name])
|
||||
.find(
|
||||
(candidate) =>
|
||||
candidate?.dims.length === 3 &&
|
||||
(candidate.type === 'float32' ||
|
||||
candidate.type === 'float64')
|
||||
)
|
||||
if (!output) {
|
||||
throw new Error('ONNX model did not return a rank-3 float output')
|
||||
}
|
||||
const [batchSize, sequenceLength, dimensions] = output.dims
|
||||
if (
|
||||
batchSize !== options.batchSize ||
|
||||
sequenceLength !== options.sequenceLength ||
|
||||
dimensions !== EMBEDDING_DIMENSIONS ||
|
||||
output.data.length !== batchSize * sequenceLength * dimensions
|
||||
) {
|
||||
throw new Error('ONNX embedding output has invalid dimensions')
|
||||
}
|
||||
|
||||
const vectors: number[][] = []
|
||||
for (let batchIndex = 0; batchIndex < batchSize; batchIndex += 1) {
|
||||
const offset = batchIndex * sequenceLength * dimensions
|
||||
const vector = Array.from(
|
||||
{ length: dimensions },
|
||||
(_, index) => output.data[offset + index] ?? Number.NaN
|
||||
)
|
||||
const magnitude = Math.hypot(...vector)
|
||||
if (!Number.isFinite(magnitude) || magnitude <= 0) {
|
||||
throw new Error('ONNX embedding output has an invalid CLS vector')
|
||||
}
|
||||
vectors.push(vector.map((component) => component / magnitude))
|
||||
}
|
||||
return vectors
|
||||
}
|
||||
|
||||
export class GraniteEmbeddingEngine implements EmbeddingInferenceEngine {
|
||||
readonly tokenizer: EmbeddingTokenizer
|
||||
readonly onnx: EmbeddingOnnxRuntime
|
||||
|
||||
private readonly dependencies: Required<GraniteEmbeddingEngineDependencies>
|
||||
private loadTask?: Promise<LoadedEngine>
|
||||
private modelDirectory?: string
|
||||
private disposed = false
|
||||
|
||||
constructor(dependencies: GraniteEmbeddingEngineDependencies = {}) {
|
||||
this.dependencies = {
|
||||
loadTokenizer:
|
||||
dependencies.loadTokenizer ??
|
||||
((tokenizerJson, tokenizerConfig) =>
|
||||
new Tokenizer(
|
||||
asJsonObject(tokenizerJson, TOKENIZER_FILE),
|
||||
asJsonObject(tokenizerConfig, TOKENIZER_CONFIG_FILE)
|
||||
)),
|
||||
createSession:
|
||||
dependencies.createSession ??
|
||||
(async (model) =>
|
||||
(await ort.InferenceSession.create(model, {
|
||||
executionProviders: ['wasm']
|
||||
})) as GraniteSessionLike),
|
||||
createInt64Tensor:
|
||||
dependencies.createInt64Tensor ??
|
||||
((data, dimensions) => new ort.Tensor('int64', data, dimensions)),
|
||||
readFile:
|
||||
dependencies.readFile ??
|
||||
(async (filePath) => new Uint8Array(await readFile(filePath)))
|
||||
}
|
||||
this.tokenizer = {
|
||||
tokenize: async ({ texts, signal }) => {
|
||||
const loaded = await this.load(this.requiredModelDirectory(), signal)
|
||||
return buildGraniteInputBatch({
|
||||
tokenizer: loaded.tokenizer,
|
||||
texts,
|
||||
padTokenId: loaded.padTokenId,
|
||||
includeTokenTypeIds:
|
||||
loaded.session.inputNames.includes('token_type_ids')
|
||||
})
|
||||
}
|
||||
}
|
||||
this.onnx = {
|
||||
run: async ({ modelDirectory, inputs, signal }) => {
|
||||
const loaded = await this.load(modelDirectory, signal)
|
||||
abortIfRequested(signal)
|
||||
const output = await loaded.session.run(
|
||||
inputs as Readonly<Record<string, unknown>>
|
||||
)
|
||||
abortIfRequested(signal)
|
||||
return output
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async embed(options: Parameters<EmbeddingInferenceEngine['embed']>[0]) {
|
||||
void options.role
|
||||
const loaded = await this.load(options.modelDirectory, options.signal)
|
||||
abortIfRequested(options.signal)
|
||||
const batch = buildGraniteInputBatch({
|
||||
tokenizer: loaded.tokenizer,
|
||||
texts: options.texts,
|
||||
padTokenId: loaded.padTokenId,
|
||||
includeTokenTypeIds:
|
||||
loaded.session.inputNames.includes('token_type_ids')
|
||||
})
|
||||
abortIfRequested(options.signal)
|
||||
const dimensions: [number, number] = [
|
||||
batch.batchSize,
|
||||
batch.sequenceLength
|
||||
]
|
||||
const feeds: Record<string, unknown> = {
|
||||
input_ids: this.dependencies.createInt64Tensor(
|
||||
batch.inputIds,
|
||||
dimensions
|
||||
),
|
||||
attention_mask: this.dependencies.createInt64Tensor(
|
||||
batch.attentionMask,
|
||||
dimensions
|
||||
)
|
||||
}
|
||||
if (batch.tokenTypeIds) {
|
||||
feeds.token_type_ids = this.dependencies.createInt64Tensor(
|
||||
batch.tokenTypeIds,
|
||||
dimensions
|
||||
)
|
||||
}
|
||||
const outputs = await loaded.session.run(feeds)
|
||||
abortIfRequested(options.signal)
|
||||
return poolGraniteClsEmbeddings({
|
||||
outputs,
|
||||
outputNames: loaded.session.outputNames,
|
||||
batchSize: batch.batchSize,
|
||||
sequenceLength: batch.sequenceLength
|
||||
})
|
||||
}
|
||||
|
||||
async dispose(): Promise<void> {
|
||||
if (this.disposed) {
|
||||
return
|
||||
}
|
||||
this.disposed = true
|
||||
const loaded = await this.loadTask?.catch(() => undefined)
|
||||
loaded?.session.release?.()
|
||||
}
|
||||
|
||||
private requiredModelDirectory(): string {
|
||||
if (!this.modelDirectory) {
|
||||
throw new Error('The Granite model has not been selected')
|
||||
}
|
||||
return this.modelDirectory
|
||||
}
|
||||
|
||||
private load(
|
||||
modelDirectory: string,
|
||||
signal: AbortSignal
|
||||
): Promise<LoadedEngine> {
|
||||
abortIfRequested(signal)
|
||||
if (this.disposed) {
|
||||
throw new Error('Granite embedding engine is disposed')
|
||||
}
|
||||
const normalizedDirectory = path.normalize(modelDirectory)
|
||||
if (
|
||||
this.modelDirectory &&
|
||||
this.modelDirectory !== normalizedDirectory
|
||||
) {
|
||||
throw new Error('Granite embedding engine cannot change model directory')
|
||||
}
|
||||
this.modelDirectory = normalizedDirectory
|
||||
this.loadTask ??= this.loadFiles(normalizedDirectory)
|
||||
return this.loadTask.then((loaded) => {
|
||||
abortIfRequested(signal)
|
||||
return loaded
|
||||
})
|
||||
}
|
||||
|
||||
private async loadFiles(modelDirectory: string): Promise<LoadedEngine> {
|
||||
const [tokenizerBytes, tokenizerConfigBytes, modelBytes] =
|
||||
await Promise.all([
|
||||
this.dependencies.readFile(
|
||||
path.join(modelDirectory, TOKENIZER_FILE)
|
||||
),
|
||||
this.dependencies.readFile(
|
||||
path.join(modelDirectory, TOKENIZER_CONFIG_FILE)
|
||||
),
|
||||
this.dependencies.readFile(path.join(modelDirectory, MODEL_FILE))
|
||||
])
|
||||
const tokenizerJson = parseJson(tokenizerBytes, TOKENIZER_FILE)
|
||||
const tokenizerConfig = parseJson(
|
||||
tokenizerConfigBytes,
|
||||
TOKENIZER_CONFIG_FILE
|
||||
)
|
||||
const tokenizer = this.dependencies.loadTokenizer(
|
||||
tokenizerJson,
|
||||
tokenizerConfig
|
||||
)
|
||||
const [padTokenId, session] = await Promise.all([
|
||||
Promise.resolve(resolvePadTokenId(tokenizer, tokenizerConfig)),
|
||||
this.dependencies.createSession(modelBytes)
|
||||
])
|
||||
return { tokenizer, padTokenId, session }
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user