Upload 3 files

MiniLMv6.cs

@@ -3,6 +3,7 @@ using UnityEngine;
 using Unity.Sentis;
 using System.IO;
 using System.Text;
+using FF = Unity.Sentis.Functional;
 
 /*
  *              Tiny Stories Inference Code
@@ -35,40 +36,43 @@ public class MiniLM : MonoBehaviour
     const int START_TOKEN = 101; 
     const int END_TOKEN = 102; 
 
-    Ops ops;
-    ITensorAllocator allocator;
-
     //Store the vocabulary
     string[] tokens;
 
-    IWorker engine;
+    const int FEATURES = 384; //size of feature space
+
+    IWorker engine, dotScore;
 
     void Start()
     {
-        allocator = new TensorCachingAllocator();
-        ops = WorkerFactory.CreateOps(backend, allocator);
-
         tokens = File.ReadAllLines(Application.streamingAssetsPath + "/vocab.txt");
 
-        Model model = ModelLoader.Load(Application.streamingAssetsPath + "/MiniLMv6.sentis");
+        engine = CreateMLModel();
 
-        engine = WorkerFactory.CreateWorker(backend, model);
+        dotScore = CreateDotScoreModel();
 
         var tokens1 = GetTokens(string1);
         var tokens2 = GetTokens(string2);
 
-        TensorFloat embedding1 = GetEmbedding(tokens1);
-        TensorFloat embedding2 = GetEmbedding(tokens2);
+        using TensorFloat embedding1 = GetEmbedding(tokens1);
+        using TensorFloat embedding2 = GetEmbedding(tokens2);
+
+        float score = GetDotScore(embedding1, embedding2);
 
-        Debug.Log("Similarity Score: " + DotScore(embedding1, embedding2));
+        Debug.Log("Similarity Score: " + score);
     }
 
-    float DotScore(TensorFloat embedding1, TensorFloat embedding2)
+    float GetDotScore(TensorFloat A, TensorFloat B)
     {
-        using var prod = ops.Mul(embedding1, embedding2);
-        using var dot = ops.ReduceSum(prod, new int[] { 1 }, false);
-        dot.MakeReadable();
-        return dot[0];
+        var inputs = new Dictionary<string, Tensor>()
+        {
+            { "input_0", A },
+            { "input_1", B }
+        };
+        dotScore.Execute(inputs);
+        var output = dotScore.PeekOutput() as TensorFloat;
+        output.CompleteOperationsAndDownload();
+        return output[0];
     }
 
     TensorFloat GetEmbedding(List<int> tokens)
@@ -85,31 +89,52 @@ public class MiniLM : MonoBehaviour
 
         var inputs = new Dictionary<string, Tensor>
         {
-            {"input_ids",input_ids },
-            {"token_type_ids",  token_type_ids},
-            {"attention_mask", attention_mask }
+            {"input_0", input_ids },
+            {"input_1", attention_mask },
+            {"input_2", token_type_ids}
         };
 
         engine.Execute(inputs);
 
-        var tokenEmbeddings = engine.PeekOutput("output") as TensorFloat;
+        var output = engine.TakeOutputOwnership("output_0") as TensorFloat;
+        return output;
+    }
+
+    IWorker CreateMLModel()
+    {
+        Model model = ModelLoader.Load(Application.streamingAssetsPath + "/MiniLMv6.sentis");
+
+        Model modelWithMeanPooling = Functional.Compile(
+          (input_ids, attention_mask, token_type_ids) =>
+          {
+              var tokenEmbeddings = model.Forward(input_ids, attention_mask, token_type_ids)[0];
+              return MeanPooling(tokenEmbeddings, attention_mask);
+          },
+          (model.inputs[0], model.inputs[1], model.inputs[2])
+        );
 
-        return MeanPooling(tokenEmbeddings, attention_mask);
+        return WorkerFactory.CreateWorker(backend, modelWithMeanPooling);
     }
 
     //Get average of token embeddings taking into account the attention mask
-    TensorFloat MeanPooling(TensorFloat tokenEmbeddings, TensorInt attentonMask)
+    FunctionalTensor MeanPooling(FunctionalTensor tokenEmbeddings, FunctionalTensor attentionMask)
     {
-        using var mask0 = attentonMask.ShallowReshape(attentonMask.shape.Unsqueeze(-1)) as TensorInt;
-        using var maskExpanded = ops.Expand(mask0, tokenEmbeddings.shape);
-        using var maskExpandedF = ops.Cast(maskExpanded, DataType.Float) as TensorFloat;
-        using var D = ops.Mul(tokenEmbeddings, maskExpandedF);
-        using var A = ops.ReduceSum(D, new[] { 1 }, false);
-        using var C = ops.ReduceSum(maskExpandedF, new[] { 1 }, false);
-        using var B = ops.Clip(C, 1e-9f, float.MaxValue);
-        using var E = ops.Div(A, B);
-        using var F = ops.ReduceL2(E, new[] { 1 }, true);
-        return ops.Div(E, F);
+        var mask = attentionMask.Unsqueeze(-1).BroadcastTo(new[] { FEATURES });     //shape=(1,N,FEATURES)
+        var A = FF.ReduceSum(tokenEmbeddings * mask, 1, false);                     //shape=(1,FEATURES)       
+        var B = A / (FF.ReduceSum(mask, 1, false) + 1e-9f);                         //shape=(1,FEATURES)
+        var C = FF.Sqrt(FF.ReduceSum(FF.Square(B), 1, true));                       //shape=(1,FEATURES)
+        return B / C;                                                               //shape=(1,FEATURES)
+    }
+
+    IWorker CreateDotScoreModel()
+    {
+        Model dotScoreModel = Functional.Compile(
+            (input1, input2) => Functional.ReduceSum(input1 * input2, 1),
+            (InputDef.Float(new TensorShape(1, FEATURES)),
+            InputDef.Float(new TensorShape(1, FEATURES)))
+        );
+
+        return WorkerFactory.CreateWorker(backend, dotScoreModel);
     }
 
     List<int> GetTokens(string text)
@@ -150,9 +175,9 @@ public class MiniLM : MonoBehaviour
     }
 
     private void OnDestroy()
-    {
+    { 
+        dotScore?.Dispose();
         engine?.Dispose();
-        ops?.Dispose();
-        allocator?.Dispose();
     }
+
 }

MiniLMv6.sentis

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cd3cc73a83d426dd085c1839e587b6a7155ce91d6698f7ae2596a3f3cd02d1cf
-size 90952597
+oid sha256:c9a2597ce9edce4c09b32e993b7f906cce91fceb2f461a597b974f71ee70453d
+size 90898400

info.json

@@ -8,7 +8,7 @@
     "data": [
         "vocab.txt"
     ],
-    "version":[
-        "1.3.0-pre.3"
+    "version": [
+        "1.4.0"
     ]
 }