LanguageModel: Require explicitly specifying the model name used

Besides providing a model name, the user can also pass an object containing the URL to a custom model. In both cases, they're explicit about the model they're exploring. As suggested by @shiffman

Author: gohai

examples/LanguageModel/sketch.js

@@ -7,7 +7,7 @@ function setup() {
   createCanvas(400, 400);
   background(0);
 
-  lm = ml5.languageModel(onModelLoaded);
+  lm = ml5.languageModel('TinyStories-15M', onModelLoaded);
 }
 
 function draw() {

examples/LanguageModelAsync/sketch.js

@@ -3,7 +3,7 @@ let lm;
 async function setup() {
   noCanvas();
 
-  lm = await ml5.languageModel();
+  lm = await ml5.languageModel('TinyStories-15M');
   console.log('Model loaded');
 
   select('#generate').mouseClicked(generateText);

examples/LanguageModelEvent/sketch.js

@@ -7,7 +7,7 @@ function setup() {
   createCanvas(400, 400);
   background(0);
 
-  lm = ml5.languageModel(onModelLoaded);
+  lm = ml5.languageModel('TinyStories-15M', onModelLoaded);
 }
 
 function draw() {

examples/LanguageModelManual/sketch.js

@@ -4,7 +4,7 @@ let numOptions = 40;
 async function setup() {
   noCanvas();
 
-  lm = await ml5.languageModel(onModelLoaded);
+  lm = await ml5.languageModel('TinyStories-15M', onModelLoaded);
 }
 
 function draw() {

examples/LanguageModelManualAsync/sketch.js

@@ -4,7 +4,7 @@ let numOptions = 40;
 async function setup() {
   noCanvas();
 
-  lm = await ml5.languageModel();
+  lm = await ml5.languageModel('TinyStories-15M');
   console.log('Model loaded');
 
   select('#generate').mouseClicked(generateText);

src/LanguageModel/index.js

@@ -5,37 +5,47 @@
 
 import { EventEmitter } from "events";
 import callCallback from "../utils/callcallback";
-import handleArguments from "../utils/handleArguments";
 
 import Llama2 from './llama2.js';
 import Llama2Wasm from './llama2.wasm';
 import Llama2Data from './llama2.data';
 
 
 class LanguageModel extends EventEmitter {
-  constructor(optionsOrCb, cb) {
+  constructor(modelNameOrOptions, callback) {
     super();
 
     this.options = {
-      modelUrl: 'https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin',  // if set, model.bin will be preloaded from provided URL (assumed to be embedded in llama2.data if not)
+      modelUrl: '',          // if set, model.bin will be preloaded from provided URL (assumed to be embedded in llama2.data if not)
       tokenizerUrl: '',      // if set, tokenizer.bin will be preloaded from provided URL (assumed to be embedded in llama2.data if not)
       steps: 0,              // how many tokens to generate (defaults to model's maximum)
       temperature: 0.9,      // 0.0 = (deterministic) argmax sampling, 1.0 = baseline
       stopOnBosOrEos: true,  // stop when encountering beginning-of-sequence or end-of-sequence token
     };
 
     // handle arguments
-    let callback;
-    if (typeof optionsOrCb === 'function') {
-      callback = optionsOrCb;
-    } else {
-      if (typeof optionsOrCb === 'object') {
-        this.options.modelUrl = (typeof optionsOrCb.modelUrl === 'string') ? optionsOrCb.modelUrl : this.options.modelUrl;
-        this.options.tokenizerUrl = (typeof optionsOrCb.tokenizerUrl === 'string') ? optionsOrCb.tokenizerUrl : this.options.tokenizerUrl;
-      }
-      if (typeof cb === 'function') {
-        callback = cb;
+    if (typeof modelNameOrOptions === 'string') {
+      switch (modelNameOrOptions) {
+        // see https://huggingface.co/karpathy/tinyllamas for TinyStories-*
+        case 'TinyStories-15M':
+          this.options.modelUrl = 'https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin';
+          break;
+        case 'TinyStories-42M':
+          this.options.modelUrl = 'https://huggingface.co/karpathy/tinyllamas/resolve/main/stories42M.bin';
+          break;
+        case 'TinyStories-110M':
+          this.options.modelUrl = 'https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.bin';
+          break;
+        default:
+          throw 'Unrecognized model ' + modelNameOrUrl + ', try e.g. TinyStories-15M';
       }
+    } else if (typeof modelNameOrUrl === 'object') {
+        this.options.modelUrl = (typeof modelNameOrOptions.modelUrl === 'string') ? modelNameOrOptions.modelUrl : this.options.modelUrl;
+        this.options.tokenizerUrl = (typeof modelNameOrOptions.tokenizerUrl === 'string') ? modelNameOrOptions.tokenizerUrl : this.options.tokenizerUrl;
+    }
+
+    if (!this.options.modelUrl) {
+      throw 'You need to provide the name of the model to load, e.g. TinyStories-15M';
     }
 
     this.out = '';
@@ -279,9 +289,8 @@ class LanguageModel extends EventEmitter {
  * exposes LanguageModel class through function
  * @returns {Object|Promise<Boolean>} A new LanguageModel instance
  */
-const languageModel = (...inputs) => {
-  const { options = {}, callback } = handleArguments(...inputs);
-  const instance = new LanguageModel(options, callback);
+const languageModel = (modelNameOrOptions, callback) => {
+  const instance = new LanguageModel(modelNameOrOptions, callback);
   return instance;
 };