feat: closer token estimations to GPT tokenizer

Author: johannschopplich

src/index.ts

@@ -76,7 +76,7 @@ export function approximateTokenSize(input: string) {
     }
     else if (ALPHANUMERIC_RE.test(token)) {
       // Increase the average token length for alphanumeric strings
-      tokenCount += Math.ceil(token.length / 4)
+      tokenCount += Math.ceil(token.length / 5)
     }
     else if (PUNCTUATION_RE.test(token)) {
       // Punctuation is often a single token, but multiple punctuations are often split

test/index.test.ts

@@ -8,14 +8,22 @@ import {
 
 describe('token-related functions', () => {
   describe('approximateTokenSize', () => {
-    it('should approximate the token size for English text', () => {
+    it('should approximate the token size for short English text', () => {
       const input = 'Hello, world! This is a test.'
-      expect(approximateTokenSize(input)).toMatchInlineSnapshot('11')
+      // 9 tokens when using the GPT tokenizer
+      expect(approximateTokenSize(input)).toMatchInlineSnapshot('9')
+    })
+
+    it('should approximate the token size for Kafka text excerpt', () => {
+      const input = 'One morning, when Gregor Samsa woke from troubled dreams, he found himself transformed in his bed into a horrible vermin. He lay on his armour-like back, and if he lifted his head a little he could see his brown belly, slightly domed and divided by arches into stiff sections. The bedding was hardly able to cover it and seemed ready to slide off any moment. His many legs, pitifully thin compared with the size of the rest of him, waved about helplessly as he looked. “What’s happened to me?” he thought. It wasn’t a dream. His room, a proper human room although a little too small, lay peacefully between its four familiar walls. A collection of textile samples lay spread out on the table – Samsa was a travelling salesman – and above it there hung a picture that he had recently cut out of an illustrated magazine and housed in a nice, gilded frame. It showed a lady fitted out with a fur hat and fur boa who sat upright, raising a heavy fur muff that covered the whole of her lower arm towards the viewer. Gregor then turned to look out the window at the dull weather. Drops of rain could be heard hitting the pane, which made him feel quite sad. “How about if I sleep a little bit longer and forget all this nonsense”, he thought, but that was something he was unable to do because he was used to sleeping on his right, and in his present state couldn’t get into that position.'
+      // 304 tokens when using the GPT tokenizer
+      expect(approximateTokenSize(input)).toMatchInlineSnapshot('364')
     })
 
     it('should approximate the token size for German text with special characters', () => {
-      const input = 'Guten Tag! Wie geht’s dir?'
-      expect(approximateTokenSize(input)).toMatchInlineSnapshot('10')
+      const input = 'Zwölf Boxkämpfer jagen Viktor quer über den großen Sylter Deich.'
+      // 22 tokens when using the GPT tokenizer
+      expect(approximateTokenSize(input)).toMatchInlineSnapshot('15')
     })
   })