Skip to content

Commit

Permalink
LanguageModel: Change the callback to be on-completion rather than on…
Browse files Browse the repository at this point in the history
…-token

As suggested by @shiffman. This makes the most basic example easier to understand.
  • Loading branch information
gohai committed Aug 4, 2023
1 parent 4eb89d9 commit b544b24
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 38 deletions.
2 changes: 1 addition & 1 deletion examples/LanguageModel/index.html
Original file line number Diff line number Diff line change
Expand Up @@ -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>
56 changes: 24 additions & 32 deletions examples/LanguageModel/sketch.js
Original file line number Diff line number Diff line change
@@ -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() {
Expand All @@ -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;
}
10 changes: 5 additions & 5 deletions src/LanguageModel/index.js
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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');

Expand Down

0 comments on commit b544b24

Please sign in to comment.