From 3ddce3c158941cb8eabdb391d207fe004095c434 Mon Sep 17 00:00:00 2001
From: Bjørn Christian Seime <bjorncs@yahooinc.com>
Date: Fri, 4 Aug 2023 13:33:13 +0200
Subject: Add generic metrics for embedders

---
 linguistics/abi-spec.json                          | 20 +++++++++++-
 .../java/com/yahoo/language/process/Embedder.java  | 37 ++++++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

(limited to 'linguistics')

diff --git a/linguistics/abi-spec.json b/linguistics/abi-spec.json
index dc85a2e6f0b..680aec3ff57 100644
--- a/linguistics/abi-spec.json
+++ b/linguistics/abi-spec.json
@@ -338,10 +338,13 @@
     ],
     "methods" : [
       "public void <init>(java.lang.String)",
+      "public com.yahoo.language.process.Embedder$Context copy()",
       "public com.yahoo.language.Language getLanguage()",
       "public com.yahoo.language.process.Embedder$Context setLanguage(com.yahoo.language.Language)",
       "public java.lang.String getDestination()",
-      "public com.yahoo.language.process.Embedder$Context setDestination(java.lang.String)"
+      "public com.yahoo.language.process.Embedder$Context setDestination(java.lang.String)",
+      "public java.lang.String getEmbedderId()",
+      "public com.yahoo.language.process.Embedder$Context setEmbedderId(java.lang.String)"
     ],
     "fields" : [ ]
   },
@@ -361,6 +364,21 @@
     ],
     "fields" : [ ]
   },
+  "com.yahoo.language.process.Embedder$Runtime" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [ ],
+    "attributes" : [
+      "public",
+      "interface",
+      "abstract"
+    ],
+    "methods" : [
+      "public abstract void sampleEmbeddingLatency(java.time.Duration, com.yahoo.language.process.Embedder$Context)",
+      "public abstract void sampleSequenceLength(long, com.yahoo.language.process.Embedder$Context)",
+      "public static com.yahoo.language.process.Embedder$Runtime testInstance()"
+    ],
+    "fields" : [ ]
+  },
   "com.yahoo.language.process.Embedder" : {
     "superClass" : "java.lang.Object",
     "interfaces" : [ ],
diff --git a/linguistics/src/main/java/com/yahoo/language/process/Embedder.java b/linguistics/src/main/java/com/yahoo/language/process/Embedder.java
index 055861c5388..1d2e1bcc847 100644
--- a/linguistics/src/main/java/com/yahoo/language/process/Embedder.java
+++ b/linguistics/src/main/java/com/yahoo/language/process/Embedder.java
@@ -1,10 +1,12 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.language.process;
 
+import com.yahoo.api.annotations.Beta;
 import com.yahoo.language.Language;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorType;
 
+import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 
@@ -64,15 +66,42 @@ public interface Embedder {
      */
     Tensor embed(String text, Context context, TensorType tensorType);
 
+    /**
+     * Runtime that is injectable through {@link Embedder} constructor.
+     */
+    @Beta
+    interface Runtime {
+        /** Sample latency metric for embedding */
+        void sampleEmbeddingLatency(Duration latency, Context ctx);
+        /** Sample sequence length metric for embedding */
+        void sampleSequenceLength(long length, Context ctx);
+
+        static Runtime testInstance() {
+            return new Runtime() {
+                @Override public void sampleEmbeddingLatency(Duration latency, Context ctx) { }
+                @Override public void sampleSequenceLength(long length, Context ctx) { }
+            };
+        }
+    }
+
     class Context {
 
         private Language language = Language.UNKNOWN;
         private String destination;
+        private String embedderId = "unknown";
 
         public Context(String destination) {
             this.destination = destination;
         }
 
+        private Context(Context other) {
+            language = other.language;
+            destination = other.destination;
+            embedderId = other.embedderId;
+        }
+
+        public Context copy() { return new Context(this); }
+
         /** Returns the language of the text, or UNKNOWN (default) to use a language independent embedding */
         public Language getLanguage() { return language; }
 
@@ -102,6 +131,14 @@ public interface Embedder {
             return this;
         }
 
+        /** Return the embedder id or 'unknown' if not set */
+        public String getEmbedderId() { return embedderId; }
+
+        /** Sets the embedder id */
+        public Context setEmbedderId(String embedderId) {
+            this.embedderId = embedderId;
+            return this;
+        }
     }
 
     class FailingEmbedder implements Embedder {
-- 
cgit v1.2.3