Add extraAbbreviations on sentence-splitter (#2029)

Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.changeset/kind-seals-bathe.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/core": patch
+---
+
+Feat: added custom abbreviations to sentence splitter

packages/core/src/llms/base.ts

@@ -1,4 +1,5 @@
-import { extractText, streamConverter } from "../utils";
+import { extractText } from "../utils/llms";
+import { streamConverter } from "../utils/stream";
 import type {
   ChatResponse,
   ChatResponseChunk,

packages/core/src/node-parser/abbreviations.ts

@@ -0,0 +1,42 @@
+export const abbreviations = {
+  english: [
+    "i.e.",
+    "etc.",
+    "vs.",
+    "Inc.",
+    "A.S.A.P.",
+    "Mr.",
+    "Mrs.",
+    "Ms.",
+    "Dr.",
+    "Prof.",
+    "Sr.",
+    "Jr.",
+    "Tel.",
+    "a.m.",
+    "p.m.",
+    "Art.",
+  ],
+  spanish: [
+    "Sr.",
+    "Sres.",
+    "Srs.",
+    "Sra.",
+    "Sras.",
+    "Srta.",
+    "Srtas.",
+    "Dr.",
+    "Drs.",
+    "Dra.",
+    "Dras.",
+    "Prof.",
+    "Profs.",
+    "Profa.",
+    "Profas.",
+    "Ing.",
+    "Lic.",
+    "Arq.",
+    "Ab.",
+    "Abs.",
+  ],
+};

packages/core/src/node-parser/sentence-splitter.ts

@@ -42,6 +42,11 @@ export class SentenceSplitter extends MetadataAwareTextSplitter {
    * Backup regex for splitting into sentences.
    */
   secondaryChunkingRegex: string = "[^,.;。？！]+[,.;。？！]?";
+  /**
+   * Extra abbreviations to consider while splitting into sentences.
+   * For example, for contracts, you may want to consider "LLC." as an important abbreviation
+   */
+  extraAbbreviations: string[] | undefined = [];
 
   #chunkingTokenizerFn = splitBySentenceTokenizer();
   #splitFns: Set<TextSplitterFn> = new Set();
@@ -59,7 +64,11 @@ export class SentenceSplitter extends MetadataAwareTextSplitter {
       this.separator = parsedParams.separator;
       this.paragraphSeparator = parsedParams.paragraphSeparator;
       this.secondaryChunkingRegex = parsedParams.secondaryChunkingRegex;
+      this.extraAbbreviations = parsedParams.extraAbbreviations;
     }
+    this.#chunkingTokenizerFn = splitBySentenceTokenizer(
+      this.extraAbbreviations,
+    );
     this.#tokenizer = params?.tokenizer ?? Settings.tokenizer;
     this.#splitFns.add(splitBySep(this.paragraphSeparator));
     this.#splitFns.add(this.#chunkingTokenizerFn);

packages/core/src/node-parser/utils.ts

@@ -1,3 +1,4 @@
+import { abbreviations } from "./abbreviations";
 import type { TextSplitter } from "./base";
 import SentenceTokenizer from "./sentence_tokenizer";
 
@@ -34,53 +35,15 @@ export const splitByChar = (): TextSplitterFn => {
   return (text: string) => text.split("");
 };
 
-let sentenceTokenizer: SentenceTokenizer | null = null;
-
-export const splitBySentenceTokenizer = (): TextSplitterFn => {
-  if (!sentenceTokenizer) {
-    sentenceTokenizer = new SentenceTokenizer([
-      // TODO: This should be improved. Take a look at: https://github.com/run-llama/LlamaIndexTS/issues/2019
-      // English
-      "i.e.",
-      "etc.",
-      "vs.",
-      "Inc.",
-      "A.S.A.P.",
-      "Mr.",
-      "Mrs.",
-      "Ms.",
-      "Dr.",
-      "Prof.",
-      "Sr.",
-      "Jr.",
-      // Spanish
-      "Sr.",
-      "Sres.",
-      "Srs.",
-      "Sra.",
-      "Sras.",
-      "Srta.",
-      "Srtas.",
-      "Dr.",
-      "Drs.",
-      "Dra.",
-      "Dras.",
-      "Prof.",
-      "Profs.",
-      "Profa.",
-      "Profas.",
-      "Ing.",
-      "Lic.",
-      "Arq.",
-      "Ab.",
-      "Abs.",
-      "Tel.",
-      "a.m.",
-      "p.m.",
-      "Art.",
-    ]);
-  }
-  const tokenizer = sentenceTokenizer;
+export const splitBySentenceTokenizer = (
+  extraAbbreviations: string[] | undefined = [],
+): TextSplitterFn => {
+  const tokenizer = new SentenceTokenizer([
+    ...abbreviations.english,
+    ...abbreviations.spanish,
+    // Add the extra abbreviations provided by the user, e.g. for business-specific context
+    ...extraAbbreviations,
+  ]);
   return (text: string) => {
     try {
       return tokenizer.tokenize(text);

packages/core/src/schema/zod.ts

@@ -51,6 +51,13 @@ export const sentenceSplitterSchema = z
       })
       .optional()
       .default("[^,.;。？！]+[,.;。？！]?"),
+    extraAbbreviations: z
+      .array(z.string(), {
+        description:
+          "Extra abbreviations to consider while splitting into sentences. For example, for contracts, you may want to consider 'LLC.' as an important abbreviation",
+      })
+      .optional()
+      .default([]),
   })
   .refine(
     (data) => data.chunkOverlap < data.chunkSize,

packages/core/src/utils/index.ts

@@ -14,19 +14,6 @@ export const isIterable = (obj: unknown): obj is Iterable<unknown> => {
   return obj != null && typeof obj === "object" && Symbol.iterator in obj;
 };
 
-export async function* streamConverter<S, D>(
-  stream: AsyncIterable<S>,
-  converter: (s: S) => D | null,
-): AsyncIterable<D> {
-  for await (const data of stream) {
-    const newData = converter(data);
-    if (newData === null) {
-      return;
-    }
-    yield newData;
-  }
-}
-
 export async function* streamCallbacks<S>(
   stream: AsyncIterable<S>,
   callbacks: {
@@ -86,3 +73,4 @@ export {
 export { MockLLM } from "./mock";
 
 export { objectEntries } from "./object-entries";
+export * from "./stream";

packages/core/src/utils/stream.ts

@@ -0,0 +1,12 @@
+export async function* streamConverter<S, D>(
+  stream: AsyncIterable<S>,
+  converter: (s: S) => D | null,
+): AsyncIterable<D> {
+  for await (const data of stream) {
+    const newData = converter(data);
+    if (newData === null) {
+      return;
+    }
+    yield newData;
+  }
+}

packages/core/tests/node-parser/text-splitter.test.ts

@@ -75,6 +75,22 @@ describe("sentence splitter", () => {
     expect(splits).toEqual(["This is a sentence. This is another sentence."]);
   });
 
+  test("overall split long text", () => {
+    const sentenceSplitter = new SentenceSplitter({
+      chunkSize: 10,
+      chunkOverlap: 0,
+    });
+    const splits = sentenceSplitter.splitText(
+      "The first short sentence. The first long long long sentence. The second short sentence. The second long long long sentence.",
+    );
+    expect(splits).toEqual([
+      "The first short sentence.",
+      "The first long long long sentence.",
+      "The second short sentence.",
+      "The second long long long sentence.",
+    ]);
+  });
+
   test("doesn't split decimals", () => {
     const sentenceSplitter = new SentenceSplitter({
       chunkSize: 5,
@@ -91,6 +107,39 @@ describe("sentence splitter", () => {
     ]);
   });
 
+  test("doesn't split basic abbreviations", () => {
+    const sentenceSplitter = new SentenceSplitter({
+      chunkSize: 15,
+      chunkOverlap: 0,
+    });
+    const splits = sentenceSplitter.splitText(
+      "This is a sentence of Broda Noel. This is the sentence of Sr. Broda Noel. This is a sentence of somebody else",
+    );
+
+    expect(splits).toEqual([
+      "This is a sentence of Broda Noel.",
+      "This is the sentence of Sr. Broda Noel.",
+      "This is a sentence of somebody else",
+    ]);
+  });
+
+  test("doesn't split extra abbreviations", () => {
+    const sentenceSplitter = new SentenceSplitter({
+      chunkSize: 10,
+      chunkOverlap: 0,
+      extraAbbreviations: ["S.A."],
+    });
+    const splits = sentenceSplitter.splitText(
+      "This is a sentence. The S.A. Broda Company. This is another sentence",
+    );
+
+    expect(splits).toEqual([
+      "This is a sentence.",
+      "The S.A. Broda Company.",
+      "This is another sentence",
+    ]);
+  });
+
   test("splits cjk", () => {
     const sentenceSplitter = new SentenceSplitter({
       chunkSize: 30,