fix: reranking edge case crash

Author: giladgd

llama/addon/AddonContext.cpp

@@ -639,6 +639,7 @@ Napi::Value AddonContext::GetEmbedding(const Napi::CallbackInfo& info) {
     }
 
     int32_t inputTokensLength = info[0].As<Napi::Number>().Int32Value();
+    int32_t maxVectorSize = (info.Length() > 1 && info[1].IsNumber()) ? info[1].As<Napi::Number>().Int32Value() : 0;
 
     if (inputTokensLength <= 0) {
         Napi::Error::New(info.Env(), "Invalid input tokens length").ThrowAsJavaScriptException();
@@ -650,15 +651,16 @@ Napi::Value AddonContext::GetEmbedding(const Napi::CallbackInfo& info) {
     const auto* embeddings = pooling_type == LLAMA_POOLING_TYPE_NONE ? NULL : llama_get_embeddings_seq(ctx, 0);
     if (embeddings == NULL) {
         embeddings = llama_get_embeddings_ith(ctx, inputTokensLength - 1);
+    }
 
-        if (embeddings == NULL) {
-            Napi::Error::New(info.Env(), std::string("Failed to get embeddings for token ") + std::to_string(inputTokensLength - 1)).ThrowAsJavaScriptException();
-            return info.Env().Undefined();
-        }
+    if (embeddings == NULL) {
+        Napi::Error::New(info.Env(), std::string("Failed to get embeddings for token ") + std::to_string(inputTokensLength - 1)).ThrowAsJavaScriptException();
+        return info.Env().Undefined();
     }
 
-    Napi::Float64Array result = Napi::Float64Array::New(info.Env(), n_embd);
-    for (size_t i = 0; i < n_embd; ++i) {
+    size_t resultSize = maxVectorSize == 0 ? n_embd : std::min(n_embd, maxVectorSize);
+    Napi::Float64Array result = Napi::Float64Array::New(info.Env(), resultSize);
+    for (size_t i = 0; i < resultSize; i++) {
         result[i] = embeddings[i];
     }
 

src/bindings/AddonTypes.ts

@@ -143,7 +143,7 @@ export type AddonContext = {
     // startPos in inclusive, endPos is exclusive
     shiftSequenceTokenCells(sequenceId: number, startPos: number, endPos: number, shiftDelta: number): void,
 
-    getEmbedding(inputTokensLength: number): Float64Array,
+    getEmbedding(inputTokensLength: number, maxVectorSize?: number): Float64Array,
     getStateSize(): number,
     getThreads(): number,
     setThreads(threads: number): void,

src/evaluator/LlamaRankingContext.ts

@@ -197,7 +197,7 @@ export class LlamaRankingContext {
                 break; // only generate one token to get embeddings
             }
 
-            const embedding = this._llamaContext._ctx.getEmbedding(input.length);
+            const embedding = this._llamaContext._ctx.getEmbedding(input.length, 1);
             if (embedding.length === 0)
                 return 0;