fix: handle sampling init failure

mybigday · Oct 20, 2023 · 43b1d71 · 43b1d71
1 parent ddd1b68
commit 43b1d71
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/android/src/main/jni.cpp b/android/src/main/jni.cpp
@@ -362,7 +362,11 @@ Java_com_rnllama_LlamaContext_doCompletion(
         env->ReleaseStringUTFChars(stop_str, stop_chars);
     }
 
-    llama->initSampling();
+    if (!llama->initSampling()) {
+        auto result = createWriteableMap(env);
+        putString(env, result, "error", "Failed to initialize sampling");
+        return reinterpret_cast<jobject>(result);
+    }
     llama->loadPrompt();
     llama->beginCompletion();
 

diff --git a/cpp/rn-llama.hpp b/cpp/rn-llama.hpp
@@ -193,11 +193,12 @@ struct llama_rn_context
         params.sparams.n_prev = n_ctx;
     }
 
-    void initSampling() {
+    bool initSampling() {
         if (ctx_sampling != nullptr) {
             llama_sampling_free(ctx_sampling);
         }
         ctx_sampling = llama_sampling_init(params.sparams);
+        return ctx_sampling != nullptr;
     }
 
     bool loadModel(gpt_params &params_)

diff --git a/ios/RNLlamaContext.mm b/ios/RNLlamaContext.mm
@@ -197,7 +197,9 @@ - (NSDictionary *)completion:(NSDictionary *)params
         }
     }
 
-    llama->initSampling();
+    if (!llama->initSampling()) {
+        @throw [NSException exceptionWithName:@"LlamaException" reason:@"Failed to initialize sampling" userInfo:nil];
+    }
     llama->loadPrompt();
     llama->beginCompletion();