fix firstToken, minor ui changes (#971)

1 files changed, 9 insertions, 2 deletions

diff --git a/candle-wasm-examples/llama2-c/llama2cWorker.js b/candle-wasm-examples/llama2-c/llama2cWorker.js
index abaf3401..a46b5bc8 100644
--- a/candle-wasm-examples/llama2-c/llama2cWorker.js
+++ b/candle-wasm-examples/llama2-c/llama2cWorker.js
@@ -50,6 +50,7 @@ async function generate(data) {
     tokenizerURL,
     prompt,
     temp,
+    top_p,
     repeatPenalty,
     seed,
     maxSeqLen,
@@ -59,11 +60,17 @@ async function generate(data) {
     const model = await Llama2C.getInstance(weightsURL, modelID, tokenizerURL);
 
     self.postMessage({ status: "loading", message: "Initializing model" });
-    model.init_with_prompt(prompt, temp, repeatPenalty, seed);
+    const firstToken = model.init_with_prompt(
+      prompt,
+      temp,
+      top_p,
+      repeatPenalty,
+      seed
+    );
 
     const seq_len = model.get_seq_len();
 
-    let sentence = "";
+    let sentence = firstToken;
     let maxTokens = maxSeqLen ? maxSeqLen : seq_len - prompt.length - 1;
     let startTime = performance.now();
     let tokensCount = 0;