Add support for (dis)allowed special and tokens in token text splitter (langchain-ai#134)

Author: dqbd

examples/package.json

@@ -19,7 +19,7 @@
   "author": "Langchain",
   "license": "MIT",
   "dependencies": {
-    "@dqbd/tiktoken": "^0.2.1",
+    "@dqbd/tiktoken": "^0.4.0",
     "chromadb": "^1.3.0",
     "langchain": "workspace:*",
     "openai": "^3.1.0",

examples/src/indexes/token_text_splitter.ts

@@ -14,6 +14,8 @@ export const run = async () => {
     encodingName: "r50k_base",
     chunkSize: 10,
     chunkOverlap: 0,
+    allowedSpecial: ["<|endoftext|>"],
+    disallowedSpecial: [],
   });
 
   const output = splitter.createDocuments([text]);

langchain/package.json

@@ -70,7 +70,7 @@
   "devDependencies": {
     "@babel/core": "^7.20.12",
     "@babel/preset-env": "^7.20.2",
-    "@dqbd/tiktoken": "^0.2.1",
+    "@dqbd/tiktoken": "^0.4.0",
     "@jest/globals": "^29.4.2",
     "@tsconfig/recommended": "^1.0.2",
     "@types/node-fetch": "2",
@@ -103,7 +103,7 @@
     "typescript": "^4.9.5"
   },
   "peerDependencies": {
-    "@dqbd/tiktoken": "^0.2.1",
+    "@dqbd/tiktoken": "^0.4.0",
     "cheerio": "^1.0.0-rc.12",
     "chromadb": "^1.3.0",
     "cohere-ai": "^5.0.2",

langchain/text_splitter.ts

@@ -185,7 +185,7 @@ export class RecursiveCharacterTextSplitter
 
 export interface TokenTextSplitterParams extends TextSplitterParams {
   encodingName: tiktoken.TiktokenEmbedding;
-  allowedSpecial: "all" | Set<string>;
+  allowedSpecial: "all" | Array<string>;
   disallowedSpecial: "all" | Array<string>;
 }
 
@@ -198,7 +198,7 @@ export class TokenTextSplitter
 {
   encodingName: tiktoken.TiktokenEmbedding;
 
-  allowedSpecial: "all" | Set<string>;
+  allowedSpecial: "all" | Array<string>;
 
   disallowedSpecial: "all" | Array<string>;
 
@@ -208,17 +208,9 @@ export class TokenTextSplitter
     super(fields);
 
     this.encodingName = fields?.encodingName ?? "gpt2";
-    this.allowedSpecial = fields?.allowedSpecial ?? new Set();
+    this.allowedSpecial = fields?.allowedSpecial ?? [];
     this.disallowedSpecial = fields?.disallowedSpecial ?? "all";
 
-    if (fields?.allowedSpecial != null) {
-      throw new Error("allowedSpecial is not implemented yet.");
-    }
-
-    if (fields?.disallowedSpecial != null) {
-      throw new Error("disallowedSpecial is not implemented yet.");
-    }
-
     try {
       const tiktoken =
         // eslint-disable-next-line @typescript-eslint/no-var-requires, global-require
@@ -235,7 +227,11 @@ export class TokenTextSplitter
   splitText(text: string): string[] {
     const splits: string[] = [];
 
-    const input_ids = this.tokenizer.encode(text);
+    const input_ids = this.tokenizer.encode(
+      text,
+      this.allowedSpecial,
+      this.disallowedSpecial
+    );
 
     let start_idx = 0;
     let cur_idx = Math.min(start_idx + this.chunkSize, input_ids.length);

yarn.lock

@@ -2248,10 +2248,10 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@dqbd/tiktoken@npm:^0.2.1":
-  version: 0.2.1
-  resolution: "@dqbd/tiktoken@npm:0.2.1"
-  checksum: 1d3fd243112b154ced97985585a439837401614e0498c9798899d925f781de769da9a0918b3f4d4c4dded1e172a546b6ec4713ad988a7990c9e1872d341b7098
+"@dqbd/tiktoken@npm:^0.4.0":
+  version: 0.4.0
+  resolution: "@dqbd/tiktoken@npm:0.4.0"
+  checksum: 2d708e70b86bd09fc9f5a5cb5e9d3bee34574e18b4c28911f311011f6673008b91510eea658e2769e29421c8e8cd7bf2c0938ffdcad58fa082871f696717383f
   languageName: node
   linkType: hard
 
@@ -9429,7 +9429,7 @@ __metadata:
   version: 0.0.0-use.local
   resolution: "langchain-examples@workspace:examples"
   dependencies:
-    "@dqbd/tiktoken": ^0.2.1
+    "@dqbd/tiktoken": ^0.4.0
     "@tsconfig/recommended": ^1.0.2
     "@typescript-eslint/eslint-plugin": ^5.51.0
     "@typescript-eslint/parser": ^5.51.0
@@ -9455,7 +9455,7 @@ __metadata:
   dependencies:
     "@babel/core": ^7.20.12
     "@babel/preset-env": ^7.20.2
-    "@dqbd/tiktoken": ^0.2.1
+    "@dqbd/tiktoken": ^0.4.0
     "@jest/globals": ^29.4.2
     "@tsconfig/recommended": ^1.0.2
     "@types/node-fetch": 2
@@ -9496,7 +9496,7 @@ __metadata:
     uuid: ^9.0.0
     yaml: ^2.2.1
   peerDependencies:
-    "@dqbd/tiktoken": ^0.2.1
+    "@dqbd/tiktoken": ^0.4.0
     cheerio: ^1.0.0-rc.12
     chromadb: ^1.3.0
     cohere-ai: ^5.0.2