Fix reasoning management in Deepseek-R1-Distill-Qwen and Qwen models

Author: orionpapadakis

src/main/java/com/example/model/qwen2/Qwen2.java

@@ -54,6 +54,32 @@ public State createNewState(int batchsize) {
         return state;
     }
 
+    /**
+     * No <|beginoftext|> needed for Qwen models.
+     */
+    @Override
+    public boolean shouldAddBeginOfText() {
+        return false;
+    }
+
+    /**
+     * No system prompt for Deepseek-R1-Distill-Qwen.
+     * Based on <a href="https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B">Usage Recommendations</a>
+     */
+    @Override
+    public boolean shouldAddSystemPrompt() {
+        return !getModelType().isDeepSeekR1();
+    }
+
+    /**
+     * Force inclusion of <think></think> for Deepseek-R1-Distill-Qwen.
+     * Based on <a href="https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B">Usage Recommendations</a>
+     */
+    @Override
+    public boolean shouldIncludeReasoning() {
+        return getModelType().isDeepSeekR1();
+    }
+
     @Override
     public void forward(State state, int token, int position) {
         if (plan == null) {

src/main/java/org/beehive/gpullama3/model/Model.java

@@ -38,6 +38,18 @@ public interface Model {
 
     State createNewState(int batchsize);
 
+    default boolean shouldAddBeginOfText() {
+        return true;
+    }
+
+    default boolean shouldAddSystemPrompt() {
+        return true;
+    }
+
+    default boolean shouldIncludeReasoning() {
+        return false;
+    }
+
     /**
      * Wrapper for invoking the model-specific forward pass via InferenceCore.
      *
@@ -68,11 +80,11 @@ default void runInteractive(Sampler sampler, Options options) {
         ChatFormat chatFormat = chatFormat();
         TornadoVMMasterPlan tornadoVMPlan = null;
 
-        if (!getModelType().equals(ModelType.QWEN_3) && !getModelType().equals(ModelType.PHI_3)) {
+        if (shouldAddBeginOfText()) {
             conversationTokens.add(chatFormat.getBeginOfText());
         }
 
-        if (options.systemPrompt() != null) {
+        if (shouldAddSystemPrompt() && options.systemPrompt() != null) {
             conversationTokens.addAll(chatFormat.encodeMessage(new ChatFormat.Message(ChatFormat.Role.SYSTEM, options.systemPrompt())));
         }
 
@@ -95,6 +107,18 @@ default void runInteractive(Sampler sampler, Options options) {
 
                 conversationTokens.addAll(chatFormat.encodeMessage(new ChatFormat.Message(ChatFormat.Role.USER, userText)));
                 conversationTokens.addAll(chatFormat.encodeHeader(new ChatFormat.Message(ChatFormat.Role.ASSISTANT, "")));
+
+                // Include reasoning for Deepseek-R1-Distill-Qwen
+                if (shouldIncludeReasoning()) {
+                    List<Integer> thinkStartTokens = tokenizer().encode("<think>\n", tokenizer().getSpecialTokens().keySet());
+                    conversationTokens.addAll(thinkStartTokens);
+
+                    // If streaming, immediately output the think start
+                    if (options.stream()) {
+                        System.out.print("<think>\n");
+                    }
+                }
+
                 Set<Integer> stopTokens = chatFormat.getStopTokens();
 
                 List<Integer> responseTokens;
@@ -127,6 +151,10 @@ default void runInteractive(Sampler sampler, Options options) {
                 }
                 if (!options.stream()) {
                     String responseText = tokenizer().decode(responseTokens);
+                    // Add the forced <think>\n prefix for non-streaming output
+                    if (shouldIncludeReasoning()) {
+                        responseText = "<think>\n" + responseText;
+                    }
                     System.out.println(responseText);
                 }
                 if (stopToken == null) {
@@ -164,11 +192,11 @@ default void runInstructOnce(Sampler sampler, Options options) {
 
         List<Integer> promptTokens = new ArrayList<>();
 
-        if (!getModelType().equals(ModelType.QWEN_3) && !getModelType().equals(ModelType.QWEN_2) && !getModelType().equals(ModelType.PHI_3)) {
+        if (shouldAddBeginOfText()) {
             promptTokens.add(chatFormat.getBeginOfText());
         }
 
-        if (options.systemPrompt() != null) {
+        if (shouldAddSystemPrompt() && options.systemPrompt() != null) {
             promptTokens.addAll(chatFormat.encodeMessage(new ChatFormat.Message(ChatFormat.Role.SYSTEM, options.systemPrompt())));
         }
 
@@ -180,6 +208,17 @@ default void runInstructOnce(Sampler sampler, Options options) {
         promptTokens.addAll(chatFormat.encodeMessage(new ChatFormat.Message(ChatFormat.Role.USER, options.prompt())));
         promptTokens.addAll(chatFormat.encodeHeader(new ChatFormat.Message(ChatFormat.Role.ASSISTANT, "")));
 
+        // Include reasoning for Deepseek-R1-Distill-Qwen
+        if (shouldIncludeReasoning()) {
+            List<Integer> thinkStartTokens = tokenizer().encode("<think>\n", tokenizer().getSpecialTokens().keySet());
+            promptTokens.addAll(thinkStartTokens);
+
+            // If streaming, immediately output the think start
+            if (options.stream()) {
+                System.out.print("<think>\n");
+            }
+        }
+
         List<Integer> responseTokens;
 
         IntConsumer tokenConsumer = token -> {
@@ -206,6 +245,10 @@ default void runInstructOnce(Sampler sampler, Options options) {
         }
         if (!options.stream()) {
             String responseText = tokenizer().decode(responseTokens);
+            // Add the forced <think>\n prefix for non-streaming output
+            if (shouldIncludeReasoning()) {
+                responseText = "<think>\n" + responseText;
+            }
             System.out.println(responseText);
         }
 

src/main/java/org/beehive/gpullama3/model/phi3/Phi3.java

@@ -53,6 +53,14 @@ public State createNewState(int batchsize) {
         return state;
     }
 
+    /**
+     * No begin of text needed for Phi3 models.
+     */
+    @Override
+    public boolean shouldAddBeginOfText() {
+        return false;
+    }
+
     @Override
     public void forward(State state, int token, int position) {
         if (plan == null) {

src/main/java/org/beehive/gpullama3/model/qwen3/Qwen3.java

@@ -53,6 +53,14 @@ public State createNewState(int batchsize) {
         return state;
     }
 
+    /**
+     * No begin of text needed for Qwen models.
+     */
+    @Override
+    public boolean shouldAddBeginOfText() {
+        return false;
+    }
+
     @Override
     public void forward(State state, int token, int position) {
         if (plan == null) {

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

@@ -53,8 +53,8 @@ public boolean isSpecialToken(int tokenIndex) {
     @Override
     public boolean shouldDisplayToken(int token) {
         int tokenType = getTokenType(token);
-
-        return tokenType == 1 || tokenType == 6;
+        // tokenType 4 allows the display of reasoning ( <think> ... <\think> )
+        return tokenType == 1 || tokenType == 4 || tokenType == 6;
     }
 
     public int getTokenType(int tokenIndex) {