update to inference engine

README.md

@@ -2,42 +2,28 @@
 license: mit
 library_name: unity-sentis
 pipeline_tag: text-generation
+tags:
+  - unity-inference-engine
 ---
 
-# Tiny Stories Model in Unity Sentis Format (Sentis 1.4.0-pre.2*)
-*Version 1.3.0 Sentis files are not compatible with Sentis 1.4.0 and would need to be recreated/downloaded
-
-This is the [Tiny Stories model](https://huggingface.co/roneneldan/TinyStories-33M) checked to run on Unity 2023. Tiny Stories is a Large Language Model that was trained on children's stories and can create stories based on the first couple of sentences.
+# Tiny Stories in Unity 6 with Inference Engine
 
+This is the [Tiny Stories model](https://huggingface.co/roneneldan/TinyStories-33M) running in Unity 6 with Inference Engine. Tiny Stories is a Large Language Model that was trained on children's stories and can create stories based on the first couple of sentences.
 
 ## How to Use
-* Create a new scene in Unity 2023
-* Install `com.unity.sentis` and `com.unity.nuget.newtonsoft-json` packages
-* Add the RunTinyStories.cs file to the Main Camera
-* Put `tinystories.sentis`, `vocab.json` and `merges.txt` in the Assets/StreamingAssets folder
-* Adjust some of the variables such as the `outputText` string to set the prompt
-* Press run
-* The output will appear in the console window
-
-## Example Input
-```
-One day an alien came down from Mars. It saw a chicken
-```
-## Example Output
-```
-One day an alien came down from Mars. It saw a chicken and said, "Hello, little chicken. What are you doing here?"
-
-The chicken replied, "I'm looking for a place to stay. I'm very tired."
-
-The alien said, "You can stay here. I have a nice place for you. It's very comfortable."
 
-The chicken was so happy. She thanked the alien and said, "Thank you. I'm very comfortable here."
+* Create a new scene in Unity 6;
+* Install `com.unity.ai.inference` from the package manager;
+* Install `com.unity.nuget.newtonsoft-json` from the package manager;
+* Drag the `tinystories.onnx` asset from the `models` folder into the `Model Asset` field;
+* Drag the `vocab.json` asset from the `data` folder into the `Vocab Asset` field;
+* Drag the `merges.txt` asset from the `data` folder into the `Merges Asset` field;
 
-The alien smiled and said, "You're welcome
-```
+## Preview
+Enter play mode. If working correctly the predicted text will be logged to the console.
 
-## Unity Sentis
-Unity Sentis is the inference engine which runs on Unity 2023. More can be found about it [here](https://unity.com/products/sentis)
+## Inference Engine
+Inference Engine is a neural network inference library for Unity. Find out more [here](https://docs.unity3d.com/Packages/com.unity.ai.inference@latest).
 
 ## Disclaimer
 The model was trained on children's stories so very unlikely to produce undesirable text. As an extra precaution, we removed a few tokens from vocab.json that might not be suitable for younger audiences. The original json can be found on the Tiny Stories original page.

RunTinyStories.cs

@@ -1,33 +1,14 @@
-using System.Collections;
 using System.Collections.Generic;
 using UnityEngine;
-using Unity.Sentis;
-using System.IO;
+using Unity.InferenceEngine;
 using System.Text;
-using FF = Unity.Sentis.Functional;
-
-/*
- *              Tiny Stories Inference Code
- *              ===========================
- *  
- *  Put this script on the Main Camera
- *  
- *  In Assets/StreamingAssets put:
- *  
- *  tinystories.sentis (or put in asset folder and drag onto field)
- *  vocab.json
- *  merges.txt
- * 
- *  Install package com.unity.nuget.newtonsoft-json from packagemanger
- *  Install package com.unity.sentis
- * 
- */
-
+using FF = Unity.InferenceEngine.Functional;
 
 public class RunTinyStories : MonoBehaviour
 {
-    //Drop the tinystories.sentis or onnx file on here if using an asset:
-    //public ModelAsset asset;
+    public ModelAsset modelAsset;
+    public TextAsset vocabAsset;
+    public TextAsset mergesAsset;
     const BackendType backend = BackendType.GPUCompute;
 
     //string outputString = "Once upon a time, there were three bears";
@@ -45,22 +26,21 @@ public class RunTinyStories : MonoBehaviour
     //Store the vocabulary
     string[] tokens;
 
-    IWorker engine;
+    Worker engine;
 
-    int currentToken = 0;
+    int currentToken;
     int[] outputTokens = new int[maxTokens];
 
     // Used for special character decoding
     int[] whiteSpaceCharacters = new int[256];
     int[] encodedCharacters = new int[256];
 
-    bool runInference = false;
-
+    bool runInference;
 
     //stop after this many tokens
     const int stopAfter = 100;
 
-    int totalTokens = 0;
+    int totalTokens;
 
     string[] merges;
     Dictionary<string, int> vocab;
@@ -71,19 +51,17 @@ public class RunTinyStories : MonoBehaviour
 
         LoadVocabulary();
 
-        var model1 = ModelLoader.Load(Path.Join(Application.streamingAssetsPath , "tinystories.sentis"));
-        //var model1 = ModelLoader.Load(asset);
+        var model1 = ModelLoader.Load(modelAsset);
         //Create a new model to select the random token:
-        var model2 = FF.Compile(
-            (input, currentToken) =>
-            {
-                var row = FF.Select(model1.Forward(input)[8], 1, currentToken);
-                return FF.Multinomial(predictability * row, 1);
-            },
-            (model1.inputs[0], InputDef.Int(new TensorShape()))
-        );
 
-        engine = WorkerFactory.CreateWorker(backend, model2);
+        var graph = new FunctionalGraph();
+        var input = graph.AddInput(model1, 0);
+        var currentTokenInput = graph.AddInput<int>(new TensorShape(), "currentToken");
+        var row = FF.Select(Functional.Forward(model1, input)[0], 1, currentTokenInput);
+        var output = FF.Multinomial(predictability * row, 1);
+        var model2 = graph.Compile(output);
+
+        engine = new Worker(model2, backend);
 
         DecodePrompt(outputString);
 
@@ -101,16 +79,14 @@ public class RunTinyStories : MonoBehaviour
 
     void RunInference()
     {
-        using var tokensSoFar = new TensorInt(new TensorShape(1, maxTokens), outputTokens);
-        using var index = new TensorInt(currentToken);
+        using var tokensSoFar = new Tensor<int>(new TensorShape(1, maxTokens), outputTokens);
+        using var index = new Tensor<int>(new TensorShape(), new[] { currentToken });
 
-        engine.Execute(new Dictionary<string, Tensor> { {"input_0", tokensSoFar },  { "input_1", index }});
+        engine.Schedule(tokensSoFar, index);
 
-        var probs = engine.PeekOutput() as TensorInt;
+        using var probs = (engine.PeekOutput() as Tensor<int>).ReadbackAndClone();
         Debug.Log(probs.shape);
 
-        probs.CompleteOperationsAndDownload();
-
         int ID = probs[0];
 
         //shift window down if got to the end
@@ -130,23 +106,22 @@ public class RunTinyStories : MonoBehaviour
         else outputString += GetUnicodeText(tokens[ID]);
 
         Debug.Log(outputString);
-
     }
 
     void DecodePrompt(string text)
     {
         var inputTokens = GetTokens(text);
 
-        for(int i = 0; i < inputTokens.Count; i++)
+        for (int i = 0; i < inputTokens.Count; i++)
         {
             outputTokens[i] = inputTokens[i];
         }
         currentToken = inputTokens.Count - 1;
     }
-   
+
     void LoadVocabulary()
     {
-        var jsonText = File.ReadAllText(Path.Join(Application.streamingAssetsPath , "vocab.json"));
+        var jsonText = vocabAsset.text;
         vocab = Newtonsoft.Json.JsonConvert.DeserializeObject<Dictionary<string, int>>(jsonText);
         tokens = new string[vocab.Count];
         foreach (var item in vocab)
@@ -154,7 +129,7 @@ public class RunTinyStories : MonoBehaviour
             tokens[item.Value] = item.Key;
         }
 
-        merges = File.ReadAllLines(Path.Join(Application.streamingAssetsPath , "merges.txt"));
+        merges = mergesAsset.text.Split("\r\n");
     }
 
     // Translates encoded special characters to Unicode
@@ -174,8 +149,7 @@ public class RunTinyStories : MonoBehaviour
         string outText = "";
         foreach (char letter in text)
         {
-            outText += ((int)letter <= 256) ? letter :
-                (char)whiteSpaceCharacters[(int)(letter - 256)];
+            outText += (letter <= 256) ? letter : (char)whiteSpaceCharacters[letter - 256];
         }
         return outText;
     }
@@ -185,7 +159,7 @@ public class RunTinyStories : MonoBehaviour
         string outText = "";
         foreach (char letter in text)
         {
-            outText += (char)encodedCharacters[(int)letter];
+            outText += (char)encodedCharacters[letter];
         }
         return outText;
     }
@@ -215,7 +189,7 @@ public class RunTinyStories : MonoBehaviour
 
         // Start with a list of single characters
         var inputTokens = new List<string>();
-        foreach(var letter in text)
+        foreach (var letter in text)
         {
             inputTokens.Add(letter.ToString());
         }
@@ -224,7 +198,7 @@ public class RunTinyStories : MonoBehaviour
 
         //Find the ids of the words in the vocab
         var ids = new List<int>();
-        foreach(var token in inputTokens)
+        foreach (var token in inputTokens)
         {
             if (vocab.TryGetValue(token, out int id))
             {
@@ -237,7 +211,7 @@ public class RunTinyStories : MonoBehaviour
 
     void ApplyMerges(List<string> inputTokens)
     {
-        foreach(var merge in merges)
+        foreach (var merge in merges)
         {
             string[] pair = merge.Split(' ');
             int n = 0;
@@ -254,9 +228,8 @@ public class RunTinyStories : MonoBehaviour
         }
     }
 
-    private void OnDestroy()
+    void OnDestroy()
     {
         engine?.Dispose();
     }
-    
 }

info.json

@@ -3,13 +3,13 @@
       "RunTinyStories.cs"
     ],
     "models": [
-        "tinystories.sentis"
+        "models/tinystories.onnx"
     ],
     "data": [
-        "vocab.json",
-        "merges.txt"
+        "data/vocab.json",
+        "data/merges.txt"
     ],
     "version": [
-        "1.4.0"
+        "2.2.0"
     ]
 }

tinystories.sentis

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c7962eb7db56b241cc19cd3f0cffcf5d76d3c35639917f07effa6b3c242c91e9
-size 478818076