Merge pull request #27387 from vespa-engine/bjorncs/hfv8.176.13

Prefer truncation configuration from tokenizer model

1 files changed, 41 insertions, 0 deletions

diff --git a/linguistics-components/src/main/java/com/yahoo/language/huggingface/ModelInfo.java b/linguistics-components/src/main/java/com/yahoo/language/huggingface/ModelInfo.java
new file mode 100644
index 00000000000..4b30b1f0435
--- /dev/null
+++ b/linguistics-components/src/main/java/com/yahoo/language/huggingface/ModelInfo.java
@@ -0,0 +1,41 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+package com.yahoo.language.huggingface;
+
+import java.util.Arrays;
+
+/**
+ * @author bjorncs
+ */
+public record ModelInfo(
+        TruncationStrategy truncation, PaddingStrategy padding, int maxLength, int stride, int padToMultipleOf) {
+
+    public enum TruncationStrategy {
+        LONGEST_FIRST,
+        ONLY_FIRST,
+        ONLY_SECOND,
+        DO_NOT_TRUNCATE;
+
+        public static TruncationStrategy fromString(String v) {
+            if ("true".equals(v)) return LONGEST_FIRST;
+            else if ("false".equals(v)) return DO_NOT_TRUNCATE;
+            return Arrays.stream(values())
+                    .filter(s -> s.name().equalsIgnoreCase(v))
+                    .findAny().orElseThrow(() -> new IllegalArgumentException("Invalid strategy '%s'".formatted(v)));
+        }
+    }
+
+    public enum PaddingStrategy {
+        LONGEST,
+        MAX_LENGTH,
+        DO_NOT_PAD;
+
+        public static PaddingStrategy fromString(String v) {
+            if ("true".equals(v)) return LONGEST;
+            else if ("false".equals(v)) return DO_NOT_PAD;
+            return Arrays.stream(values())
+                    .filter(s -> s.name().equalsIgnoreCase(v))
+                    .findAny().orElseThrow(() -> new IllegalArgumentException("Invalid strategy '%s'".formatted(v)));
+        }
+    }
+}