Merge pull request #124 from jiaw3i/main

Author: PlexPt

pom.xml

@@ -135,7 +135,11 @@
             <version>4.13.2</version>
             <scope>test</scope>
         </dependency>
-
+        <dependency>
+            <groupId>com.knuddels</groupId>
+            <artifactId>jtokkit</artifactId>
+            <version>0.4.0</version>
+        </dependency>
     </dependencies>
 
     <!-- 下面这个标签里的不能改 -->

src/main/java/com/plexpt/chatgpt/entity/chat/ChatCompletion.java

@@ -14,6 +14,7 @@
 import lombok.NoArgsConstructor;
 import lombok.NonNull;
 import lombok.extern.slf4j.Slf4j;
+import com.plexpt.chatgpt.util.TokensUtil;
 
 /**
  * chat
@@ -24,7 +25,7 @@
 @Builder
 @Slf4j
 @AllArgsConstructor
-@NoArgsConstructor
+@NoArgsConstructor(force = true)
 @JsonInclude(JsonInclude.Include.NON_NULL)
 public class ChatCompletion implements Serializable {
 
@@ -125,6 +126,9 @@ public enum Model {
         private String name;
     }
 
+    public int countTokens() {
+        return TokensUtil.tokens(this.model, this.messages);
+    }
 }
 
 

src/main/java/com/plexpt/chatgpt/entity/chat/Message.java

@@ -1,5 +1,6 @@
 package com.plexpt.chatgpt.entity.chat;
 
+import com.fasterxml.jackson.annotation.JsonInclude;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.Data;
@@ -13,12 +14,19 @@
 @AllArgsConstructor
 @NoArgsConstructor
 @Builder
+@JsonInclude(JsonInclude.Include.NON_NULL)
 public class Message {
     /**
      * 目前支持三种角色参考官网，进行情景输入：https://platform.openai.com/docs/guides/chat/introduction
      */
     private String role;
     private String content;
+    private String name;
+
+    public Message(String role, String content) {
+        this.role = role;
+        this.content = content;
+    }
 
     public static Message of(String content) {
 

src/main/java/com/plexpt/chatgpt/util/TokensUtil.java

@@ -0,0 +1,64 @@
+package com.plexpt.chatgpt.util;
+
+import cn.hutool.core.util.StrUtil;
+import com.knuddels.jtokkit.Encodings;
+import com.knuddels.jtokkit.api.Encoding;
+import com.knuddels.jtokkit.api.EncodingRegistry;
+import com.plexpt.chatgpt.entity.chat.ChatCompletion;
+import com.plexpt.chatgpt.entity.chat.Message;
+import lombok.experimental.UtilityClass;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+@UtilityClass
+public class TokensUtil {
+
+    private static final Map<String, Encoding> modelEncodingMap = new HashMap<>();
+    private static final EncodingRegistry encodingRegistry = Encodings.newDefaultEncodingRegistry();
+
+    static {
+        for (ChatCompletion.Model model : ChatCompletion.Model.values()) {
+            Optional<Encoding> encodingForModel = encodingRegistry.getEncodingForModel(model.getName());
+            encodingForModel.ifPresent(encoding -> modelEncodingMap.put(model.getName(), encoding));
+        }
+    }
+
+    /**
+     * 计算tokens
+     * @param modelName 模型名称
+     * @param messages 消息列表
+     * @return 计算出的tokens数量
+     */
+
+    public static int tokens(String modelName, List<Message> messages) {
+        Encoding encoding = modelEncodingMap.get(modelName);
+        if (encoding == null) {
+            throw new IllegalArgumentException("Unsupported model: " + modelName);
+        }
+
+        int tokensPerMessage = 0;
+        int tokensPerName = 0;
+        if (modelName.startsWith("gpt-4")) {
+            tokensPerMessage = 3;
+            tokensPerName = 1;
+        } else if (modelName.startsWith("gpt-3.5-turbo")) {
+            tokensPerMessage = 4; // every message follows <|start|>{role/name}\n{content}<|end|>\n
+            tokensPerName = -1; // if there's a name, the role is omitted
+        }
+        int sum = 0;
+        for (Message message : messages) {
+            sum += tokensPerMessage;
+            sum += encoding.countTokens(message.getContent());
+            sum += encoding.countTokens(message.getRole());
+            if (StrUtil.isNotBlank(message.getName())) {
+                sum += encoding.countTokens(message.getName());
+                sum += tokensPerName;
+            }
+        }
+        sum += 3;
+        return sum;
+    }
+}

src/test/java/com/plexpt/chatgpt/Test.java

@@ -61,4 +61,29 @@ public void chatmsg() {
         System.out.println(res);
     }
 
+    /**
+     * 测试tokens数量计算
+     */
+    @org.junit.Test
+    public void tokens() {
+        Message system = Message.ofSystem("你现在是一个诗人，专门写七言绝句");
+        Message message = Message.of("写一段七言绝句诗，题目是：火锅！");
+
+        ChatCompletion chatCompletion1 = ChatCompletion.builder()
+                .model(ChatCompletion.Model.GPT_3_5_TURBO.getName())
+                .messages(Arrays.asList(system, message))
+                .maxTokens(3000)
+                .temperature(0.9)
+                .build();
+        ChatCompletion chatCompletion2 = ChatCompletion.builder()
+                .model(ChatCompletion.Model.GPT_4.getName())
+                .messages(Arrays.asList(system, message))
+                .maxTokens(3000)
+                .temperature(0.9)
+                .build();
+
+        log.info("{} tokens: {}", chatCompletion1.getModel(), chatCompletion1.countTokens());
+        log.info("{} tokens: {}", chatCompletion2.getModel(), chatCompletion2.countTokens());
+    }
+
 }