diff --git a/src/evaluator/LlamaEmbeddingContext.ts b/src/evaluator/LlamaEmbeddingContext.ts index b814b034..792deabe 100644 --- a/src/evaluator/LlamaEmbeddingContext.ts +++ b/src/evaluator/LlamaEmbeddingContext.ts @@ -91,6 +91,15 @@ export class LlamaEmbeddingContext { "Try to increase the context size or use another model that supports longer contexts." ); + if ( + this.model.fileInfo.architectureMetadata.attention?.causal === false && + resolvedInput.length > this._llamaContext.batchSize + ) + throw new Error( + "Input is longer than the batch size for this non-causal embedding model. " + + "Try to increase the batch size to at least the input length." + ); + return await withLock([this as LlamaEmbeddingContext, "evaluate"], async () => { await this._sequence.eraseContextTokenRanges([{ start: 0, diff --git a/test/modelDependent/nomicEmbedText/embedding.test.ts b/test/modelDependent/nomicEmbedText/embedding.test.ts index a0c4848b..c916b247 100644 --- a/test/modelDependent/nomicEmbedText/embedding.test.ts +++ b/test/modelDependent/nomicEmbedText/embedding.test.ts @@ -57,6 +57,24 @@ describe("nomic embed text", () => { expect(helloThereEmbedding2.vector).to.eql(helloThereEmbedding.vector); }); + test("rejects non-causal inputs longer than the batch size", {timeout: 1000 * 60 * 60 * 2}, async () => { + const modelPath = await getModelFile("nomic-embed-text-v1.5.Q4_K_M.gguf"); + const llama = await getTestLlama(); + + const model = await llama.loadModel({ + modelPath + }); + const embeddingContext = await model.createEmbeddingContext({ + contextSize: 2048 + }); + const input = "word ".repeat(600); + + expect(model.tokenize(input).length).toBeGreaterThan(embeddingContext._llamaContext.batchSize); + await expect(embeddingContext.getEmbeddingFor(input)).rejects.toThrow( + "Input is longer than the batch size for this non-causal embedding model." + ); + }); + test("similarity search", {timeout: 1000 * 60 * 60 * 2}, async () => { const modelPath = await getModelFile("nomic-embed-text-v1.5.Q4_K_M.gguf"); const llama = await getTestLlama();