Merge pull request #53 from beehive-lab/demo-integration

Add support for encoding ordinary text in Qwen3Tokenizer and update Q…

Author: mikepapadim

src/main/java/org/beehive/gpullama3/model/format/Qwen3ChatFormat.java

@@ -2,10 +2,7 @@
 
 import org.beehive.gpullama3.tokenizer.impl.Qwen3Tokenizer;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
 /**
  * Utility tailored for the Chat Markup Language (ChatML) prompt format.
@@ -42,9 +39,9 @@ public Qwen3ChatFormat(Qwen3Tokenizer tokenizer, ChatTokens chatTokens) {
         this.imStart = startHeader;
         this.imEnd = endHeader;
 
-        fimPrefix = specialTokens.getOrDefault("<|fim_prefix|>", -1);
-        fimSuffix = specialTokens.getOrDefault("<|fim_suffix|>", -1);
-        fimMiddle = specialTokens.getOrDefault("<|fim_middle|>", -1);
+        this.fimPrefix = specialTokens.getOrDefault("<|fim_prefix|>", -1);
+        this.fimSuffix = specialTokens.getOrDefault("<|fim_suffix|>", -1);
+        this.fimMiddle = specialTokens.getOrDefault("<|fim_middle|>", -1);
     }
 
     public ChatTokens chatTokens() {
@@ -66,7 +63,7 @@ public List<Integer> encodeHeader(Message message) {
                 default -> null;
             };
             if (sToken != null) {
-                Integer token = tokenizer.getSpecialTokens().get("<｜User｜>");
+                Integer token = tokenizer.getSpecialTokens().get(sToken);
                 if (token == null) {
                     throw new IllegalStateException(String.format("Unknown token '%s'", sToken));
                 }
@@ -80,19 +77,23 @@ public List<Integer> encodeHeader(Message message) {
         } else if (Role.FIM_MIDDLE.equals(message.role())) {
             tokens.add(fimMiddle);
         } else {
+            // Add the special token directly, don't try to encode it
             tokens.add(imStart);
-            tokens.addAll(this.tokenizer.encodeAsList(message.role().name()));
-            tokens.addAll(this.tokenizer.encodeAsList("\n"));
+            // Encode the role name as ordinary text (no special tokens in role names)
+            tokens.addAll(this.tokenizer.encodeOrdinaryAsList(message.role().name()));
+            tokens.addAll(this.tokenizer.encodeOrdinaryAsList("\n"));
         }
         return tokens;
     }
 
     @Override
     public List<Integer> encodeMessage(Message message) {
         List<Integer> tokens = this.encodeHeader(message);
-        tokens.addAll(this.tokenizer.encodeAsList(message.content().strip()));
+        // Encode message content as ordinary text
+        tokens.addAll(this.tokenizer.encodeOrdinaryAsList(message.content().strip()));
         boolean isFim = Role.FIM_PREFIX.equals(message.role()) || Role.FIM_SUFFIX.equals(message.role()) || Role.FIM_MIDDLE.equals(message.role());
         if (imEnd != -1 && !isFim) {
+            // Add the end token directly
             tokens.add(imEnd);
         }
         return tokens;
@@ -108,9 +109,19 @@ public Set<Integer> getStopTokens() {
         if (imEnd == -1 && endOfText == -1) {
             throw new IllegalStateException("No stop token is defined.");
         }
-        if (imEnd == -1) {
-            return Set.of(endOfText);
+
+        // Only add valid token IDs (not -1)
+        Set<Integer> stopTokens = new HashSet<>();
+        if (imEnd != -1) {
+            stopTokens.add(imEnd);
+        }
+        if (endOfText != -1) {
+            stopTokens.add(endOfText);
         }
-        return Set.of(imEnd, endOfText, endOfTextFim);
+        if (endOfTextFim != -1) {
+            stopTokens.add(endOfTextFim);
+        }
+
+        return stopTokens;
     }
 }

src/main/java/org/beehive/gpullama3/tokenizer/impl/Qwen3Tokenizer.java

@@ -271,11 +271,27 @@ public List<Integer> encode(String text, Set<String> allowedSpecial) {
     }
     // @formatter:on
 
+    /**
+     * Encode text as ordinary tokens (no special token handling)
+     */
+    public List<Integer> encodeOrdinaryAsList(String text) {
+        // First convert to byte-encoded unicode representation
+        StringBuilder sb = new StringBuilder();
+        byte[] bytes = text.getBytes(StandardCharsets.UTF_8);
+        for (byte b : bytes) {
+            sb.appendCodePoint(BYTE_ENCODER.get(Byte.toUnsignedInt(b)));
+        }
+        // Then encode using BPE
+        return encodeOrdinary(sb.toString());
+    }
+
     @Override
     public List<Integer> encodeAsList(String text) {
         return Arrays.stream(encode(text)).boxed().toList();
     }
 
+
+
     public String decodeImpl(List<Integer> tokens) {
         StringBuilder sb = new StringBuilder();
         for (int token : tokens) {