Merge pull request #1 from paradite/feat/token-counter

Feat/token counter
paradite · Jul 21, 2024 · 139722d · 139722d
2 parents 2120f9a + 59bd0d9
commit 139722d
Show file tree

Hide file tree

Showing 11 changed files with 753 additions and 123 deletions.
diff --git a/README.md b/README.md
@@ -32,6 +32,7 @@ $ yarn add llm-info
 ## Usage
 
 ```ts
+// Models
 import { AllModels, ModelEnum, NonModelEnum, ModelInfoMap } from 'llm-info';
 
 console.log(AllModels);
@@ -48,45 +49,42 @@ console.log(AllModelLikes);
 ]
 */
 
-const model = ModelEnum['gpt-4'];
-const modelInfo = ModelInfoMap[model];
+// Model Info
+const modelInfo = ModelInfoMap['gpt-4o'];
 console.log(modelInfo);
 /*
 {
-  name: 'GPT-4',
+  name: 'GPT-4o',
   provider: 'openai',
   contextWindowTokenLimit: 128000,
   outputTokenLimit: 4096,
-  pricePerMillionInputTokens: 30,
-  pricePerMillionOutputTokens: 60
+  pricePerMillionInputTokens: 5,
+  pricePerMillionOutputTokens: 15,
+  tokenizerId: 'Xenova/gpt-4o'
 }
 */
 
-console.log(ModelInfoMap[ModelEnum['claude-3-5-sonnet-20240620']]);
-/*
-{
-  name: 'Claude 3.5 Sonnet',
-  provider: 'anthropic',
-  contextWindowTokenLimit: 200000,
-  outputTokenLimit: 4096,
-  pricePerMillionInputTokens: 3,
-  pricePerMillionOutputTokens: 15
+// Tokenizer
+import { AutoTokenizer } from '@xenova/transformers';
+const testSentence =
+  "Many words map to one token, but some don't: indivisible.";
+const results: string[] = [];
+for (let i = 0; i < AllModels.length; i++) {
+  const model = AllModels[i];
+  if (ModelInfoMap[model].tokenizerId) {
+    const tokenizer = await AutoTokenizer.from_pretrained(
+      ModelInfoMap[model].tokenizerId
+    );
+    const tokens = tokenizer.encode(testSentence);
+    results.push(`${model}: ${tokens.length}`);
+  }
 }
-*/
-
-const modelLike = NonModelEnum['chatgpt'];
-const modelLikeInfo = ModelInfoMap[modelLike];
-console.log(modelLikeInfo);
-/*
-{
-  name: 'ChatGPT',
-  provider: 'openai',
-  contextWindowTokenLimit: 4096,
-  outputTokenLimit: 4096,
-  pricePerMillionInputTokens: null,
-  pricePerMillionOutputTokens: null
-}
-*/
+console.log(`Test sentence: ${testSentence}\n${results.join('\n')}`);
+// Test sentence: Many words map to one token, but some don't: indivisible.
+// gpt-4: 15
+// gpt-4o: 14
+// gpt-4o-mini: 14
+// claude-3-5-sonnet-20240620: 16
 ```
 
 ## Testing

diff --git a/jest.config.js b/jest.config.js
diff --git a/jest.config.ts b/jest.config.ts
@@ -0,0 +1,24 @@
+// jest.config.ts
+import type { JestConfigWithTsJest } from 'ts-jest';
+
+// https://kulshekhar.github.io/ts-jest/docs/guides/esm-support
+
+const jestConfig: JestConfigWithTsJest = {
+  // [...]
+  extensionsToTreatAsEsm: ['.ts'],
+  moduleNameMapper: {
+    '^(\\.{1,2}/.*)\\.js$': '$1',
+  },
+  transform: {
+    // '^.+\\.[tj]sx?$' to process ts,js,tsx,jsx with `ts-jest`
+    // '^.+\\.m?[tj]sx?$' to process ts,js,tsx,jsx,mts,mjs,mtsx,mjsx with `ts-jest`
+    '^.+\\.tsx?$': [
+      'ts-jest',
+      {
+        useESM: true,
+      },
+    ],
+  },
+};
+
+export default jestConfig;
diff --git a/package.json b/package.json
@@ -1,8 +1,9 @@
 {
   "name": "llm-info",
-  "version": "1.0.3",
+  "version": "1.0.4",
   "description": "Information on LLM models, context window token limit, output token limit, pricing and more",
   "main": "dist/index.js",
+  "type": "module",
   "module": "./dist/index.mjs",
   "types": "./dist/index.d.ts",
   "exports": {
@@ -15,7 +16,7 @@
   "scripts": {
     "compile": "tsup src/index.ts --dts --format esm,cjs",
     "prepublishOnly": "rm -rf ./dist && npm run compile",
-    "test": "jest"
+    "test": "NODE_OPTIONS=\"$NODE_OPTIONS --experimental-vm-modules\" jest"
   },
   "author": "paradite",
   "license": "MIT",
@@ -31,8 +32,10 @@
   ],
   "devDependencies": {
     "@types/jest": "^29.5.12",
+    "@xenova/transformers": "^2.17.2",
     "jest": "^29.7.0",
     "ts-jest": "^29.2.3",
+    "ts-node": "^10.9.2",
     "tsup": "^8.2.0",
     "typescript": "^5.5.3"
   }

diff --git a/src/index.ts b/src/index.ts
@@ -1,75 +1,3 @@
-export enum ModelEnum {
-  'gpt-4' = 'gpt-4',
-  'gpt-4o' = 'gpt-4o',
-  'gpt-4o-mini' = 'gpt-4o-mini',
-  'claude-3-5-sonnet-20240620' = 'claude-3-5-sonnet-20240620',
-}
-
-export enum NonModelEnum {
-  'chatgpt' = 'chatgpt',
-}
-
-export const AllModels = Object.values(ModelEnum);
-
-export const AllModelLikes = [...AllModels, ...Object.values(NonModelEnum)];
-
-export type ModelLike = ModelEnum | NonModelEnum;
-
-export const AI_PROVIDERS = {
-  OPENAI: 'openai',
-  ANTHROPIC: 'anthropic',
-} as const;
-
-export type AI_PROVIDER_TYPE = (typeof AI_PROVIDERS)[keyof typeof AI_PROVIDERS];
-
-export type ModelInfo = {
-  name: string;
-  provider: AI_PROVIDER_TYPE;
-  contextWindowTokenLimit: number;
-  outputTokenLimit: number;
-  pricePerMillionInputTokens: number | null;
-  pricePerMillionOutputTokens: number | null;
-};
-
-export const ModelInfoMap: Record<ModelLike, ModelInfo> = {
-  [ModelEnum['gpt-4']]: {
-    name: 'GPT-4',
-    provider: AI_PROVIDERS.OPENAI,
-    contextWindowTokenLimit: 128000,
-    outputTokenLimit: 4096,
-    pricePerMillionInputTokens: 30,
-    pricePerMillionOutputTokens: 60,
-  },
-  [ModelEnum['gpt-4o']]: {
-    name: 'GPT-4o',
-    provider: AI_PROVIDERS.OPENAI,
-    contextWindowTokenLimit: 128000,
-    outputTokenLimit: 4096,
-    pricePerMillionInputTokens: 5,
-    pricePerMillionOutputTokens: 15,
-  },
-  [ModelEnum['gpt-4o-mini']]: {
-    name: 'GPT-4o mini',
-    provider: AI_PROVIDERS.OPENAI,
-    contextWindowTokenLimit: 128000,
-    outputTokenLimit: 4096,
-    pricePerMillionInputTokens: 0.15,
-    pricePerMillionOutputTokens: 0.6,
-  },
-  [ModelEnum['claude-3-5-sonnet-20240620']]: {
-    name: 'Claude 3.5 Sonnet',
-    provider: AI_PROVIDERS.ANTHROPIC,
-    contextWindowTokenLimit: 200000,
-    outputTokenLimit: 4096,
-    pricePerMillionInputTokens: 3,
-    pricePerMillionOutputTokens: 15,
-  },
-  [NonModelEnum['chatgpt']]: {
-    name: 'ChatGPT',
-    provider: AI_PROVIDERS.OPENAI,
-    contextWindowTokenLimit: 4096,
-    outputTokenLimit: 4096,
-    pricePerMillionInputTokens: null,
-    pricePerMillionOutputTokens: null,
-  },
-};
+export * from './model';
+export * from './provider';
+export * from './modelInfo';
diff --git a/src/model.ts b/src/model.ts
@@ -0,0 +1,16 @@
+export enum ModelEnum {
+  'gpt-4' = 'gpt-4',
+  'gpt-4o' = 'gpt-4o',
+  'gpt-4o-mini' = 'gpt-4o-mini',
+  'claude-3-5-sonnet-20240620' = 'claude-3-5-sonnet-20240620',
+}
+
+export enum NonModelEnum {
+  'chatgpt' = 'chatgpt',
+}
+
+export const AllModels = Object.values(ModelEnum);
+
+export const AllModelLikes = [...AllModels, ...Object.values(NonModelEnum)];
+
+export type ModelLike = ModelEnum | NonModelEnum;
diff --git a/src/modelInfo.ts b/src/modelInfo.ts
@@ -0,0 +1,60 @@
+import { ModelEnum, ModelLike, NonModelEnum } from './model';
+import { AI_PROVIDER_TYPE, AI_PROVIDERS } from './provider';
+
+export type ModelInfo = {
+  name: string;
+  provider: AI_PROVIDER_TYPE;
+  contextWindowTokenLimit: number;
+  outputTokenLimit: number;
+  pricePerMillionInputTokens: number | null;
+  pricePerMillionOutputTokens: number | null;
+  tokenizerId: string | null;
+};
+
+export const ModelInfoMap: Record<ModelLike, ModelInfo> = {
+  [ModelEnum['gpt-4']]: {
+    name: 'GPT-4',
+    provider: AI_PROVIDERS.OPENAI,
+    contextWindowTokenLimit: 128000,
+    outputTokenLimit: 4096,
+    pricePerMillionInputTokens: 30,
+    pricePerMillionOutputTokens: 60,
+    tokenizerId: 'Xenova/gpt-4',
+  },
+  [ModelEnum['gpt-4o']]: {
+    name: 'GPT-4o',
+    provider: AI_PROVIDERS.OPENAI,
+    contextWindowTokenLimit: 128000,
+    outputTokenLimit: 4096,
+    pricePerMillionInputTokens: 5,
+    pricePerMillionOutputTokens: 15,
+    tokenizerId: 'Xenova/gpt-4o',
+  },
+  [ModelEnum['gpt-4o-mini']]: {
+    name: 'GPT-4o mini',
+    provider: AI_PROVIDERS.OPENAI,
+    contextWindowTokenLimit: 128000,
+    outputTokenLimit: 4096,
+    pricePerMillionInputTokens: 0.15,
+    pricePerMillionOutputTokens: 0.6,
+    tokenizerId: 'Xenova/gpt-4o',
+  },
+  [ModelEnum['claude-3-5-sonnet-20240620']]: {
+    name: 'Claude 3.5 Sonnet',
+    provider: AI_PROVIDERS.ANTHROPIC,
+    contextWindowTokenLimit: 200000,
+    outputTokenLimit: 4096,
+    pricePerMillionInputTokens: 3,
+    pricePerMillionOutputTokens: 15,
+    tokenizerId: 'Xenova/claude-tokenizer',
+  },
+  [NonModelEnum['chatgpt']]: {
+    name: 'ChatGPT',
+    provider: AI_PROVIDERS.OPENAI,
+    contextWindowTokenLimit: 4096,
+    outputTokenLimit: 4096,
+    pricePerMillionInputTokens: null,
+    pricePerMillionOutputTokens: null,
+    tokenizerId: null,
+  },
+};
diff --git a/src/provider.ts b/src/provider.ts
@@ -0,0 +1,6 @@
+export const AI_PROVIDERS = {
+  OPENAI: 'openai',
+  ANTHROPIC: 'anthropic',
+} as const;
+
+export type AI_PROVIDER_TYPE = (typeof AI_PROVIDERS)[keyof typeof AI_PROVIDERS];
diff --git a/test/index.test.ts b/test/index.test.ts
@@ -26,10 +26,9 @@ describe('llm', () => {
     ]);
   });
   it('info works', () => {
-    const model = ModelEnum['gpt-4'];
-    const modelInfo = ModelInfoMap[model];
+    const modelInfo = ModelInfoMap['gpt-4o'];
     console.log(modelInfo);
-    expect(modelInfo.name).toBe('GPT-4');
+    expect(modelInfo.name).toBe('GPT-4o');
 
     console.log(ModelInfoMap[ModelEnum['claude-3-5-sonnet-20240620']]);
 

diff --git a/test/tokenizer.test.ts b/test/tokenizer.test.ts
@@ -0,0 +1,23 @@
+import { AllModels, ModelInfoMap } from '../src';
+import { AutoTokenizer } from '@xenova/transformers';
+
+describe('llm', () => {
+  it('tokenizer works', async () => {
+    const testSentence =
+      "Many words map to one token, but some don't: indivisible.";
+    const results: string[] = [];
+    for (let i = 0; i < AllModels.length; i++) {
+      const model = AllModels[i];
+      if (ModelInfoMap[model].tokenizerId) {
+        const tokenizer = await AutoTokenizer.from_pretrained(
+          ModelInfoMap[model].tokenizerId
+        );
+        const tokens = tokenizer.encode(testSentence);
+        expect(tokens.length).toBeGreaterThanOrEqual(14);
+        expect(tokens.length).toBeLessThanOrEqual(17);
+        results.push(`${model}: ${tokens.length}`);
+      }
+    }
+    console.log(`Test sentence: ${testSentence}\n${results.join('\n')}`);
+  }, 5000);
+});