LanguageModel: Change the callback to be on-completion rather than on-token

gohai · gohai · commit 8225b4c40042 · 2023-08-09T19:48:31.000+08:00
As suggested by @shiffman. This makes the most basic example easier to understand.
diff --git a/examples/LanguageModel/index.html b/examples/LanguageModel/index.html
@@ -8,7 +8,7 @@
   <script src="../../dist/ml5.js"></script>
 </head>
 <body>
-  <input placeholder="Prompt" id="prompt"></input> <input type="button" value="Generate" id="generate" disabled></input>
+  <input placeholder="Prompt" id="prompt"></input> <input type="button" value="Generate" id="generate" disabled></input><br><br>
   <script src="sketch.js"></script>
 </body>
 </html>
diff --git a/examples/LanguageModel/sketch.js b/examples/LanguageModel/sketch.js
@@ -1,18 +1,32 @@
 let lm;
+let words = [];
 let curWord = 0;
-let angle = 0;
-let radius = 60;
+let curX = 0;
+let curY = 20;
 
 function setup() {
   createCanvas(400, 400);
+  frameRate(15);
   background(0);
 
   lm = ml5.languageModel('TinyStories-15M', onModelLoaded);
 }
 
 function draw() {
-  fill(0, 5);
+  fill(0, 2);
   rect(0, 0, width, height);
+
+  if (curWord < words.length) {
+    let wordWidth = textWidth(words[curWord]);
+    if (curX+wordWidth > width) {
+      curX = 0;
+      curY += textAscent() + textDescent();
+    }
+    fill(255);
+    text(words[curWord] + ' ', curX, curY);
+    curX += wordWidth;
+    curWord++;
+  }
 }
 
 function onModelLoaded() {
@@ -23,39 +37,17 @@ function onModelLoaded() {
 
 function generateText() {
   let prompt = select('#prompt').value();
-  console.log('Prompt: ' + prompt);
+  console.log('Prompt is "' + prompt + '"');
 
   let options = {
     temperature: 0.9
   };
-  lm.generate(prompt, options, onToken);
-
-  curWord = 0;
+  lm.generate(prompt, options, gotText);
 }
 
-function onToken(lm) {
-  console.log(lm.tokens[lm.tokens.length-1]);
-
-  while (curWord < lm.words.length) {
-    push();
-    translate(width/2, height/2);
-    rotate(radians(angle));
-    translate(radius, 0);
-    rotate(radians(-angle));
-    angle = angle + 5;
-    radius += 0.5;
-    if (radius > width/2 || radius > height/2) {
-      radius = 60;
-    }
-    fill(255);
-    textAlign(CENTER);
-    textSize(10);
-    text(lm.words[curWord], 0, 0);
-    pop();
-    curWord++;
-  }
-
-  if (lm.finished) {
-    console.log('Generation finished');
-  }
+function gotText(out, lm) {
+  console.log('Model returned "' + out + '"');
+  // lm.words contains the output broken up in words
+  words = lm.words;
+  curWord = 0;
 }
diff --git a/src/LanguageModel/index.js b/src/LanguageModel/index.js
@@ -108,10 +108,7 @@ class LanguageModel extends EventEmitter {
         this.out += tokenStr;
       }
 
-      // on-token callback/event
-      if (this.callback) {
-        this.callback(this);
-      }
+      // on-token event
       this.emit('token', this);
 
       // redo word tokenization
@@ -128,13 +125,16 @@ class LanguageModel extends EventEmitter {
         this.emit('word', this.words[i], this);
       }
 
-      // on-finished promise/event
+      // on-finished promise/event/callback
       if (this.finished) {
         // fulfill the promise returned by generate()
         if (this.promiseResolve) {
           this.promiseResolve(this.out);
         }
         this.emit('finsh', this);
+        if (this.callback) {
+          this.callback(this.out, this);
+        }
       }
     }, 'viifi');
 
