gpt-token-utils
Version:
Isomorphic utilities for GPT-3 tokenization and prompt building.
49 lines (48 loc) • 1.4 kB
JavaScript
/**
* @copyright Sister Software. All rights reserved.
* @author Teffen Ellis, et al.
* @license
* See LICENSE file in the project root for full license information.
*/
/**
* The IDs of available model families.
*/
export const ModelFamilyIDs = {
Ada: 'ada',
Babbage: 'babbage',
Curie: 'curie',
Davinci: 'davinci',
ChatGPT: 'chat-gpt',
GPT4: 'gpt-4',
GPT4_32K: 'gpt-4-32k',
};
export class ModelFamiliesMap {
_familyMap = new Map();
_modelToFamilyMap = new Map();
addFamily(family) {
this._familyMap.set(family.familyID, family);
for (const modelID of family.modelIDs) {
this._modelToFamilyMap.set(modelID, family);
}
}
getFamilyByFamilyID(familyID) {
return this._familyMap.get(familyID);
}
getFamilyByModelID(modelID) {
return this._modelToFamilyMap.get(modelID);
}
get = (input) => {
if (typeof input === 'string') {
const family = this.getFamilyByFamilyID(input) || this.getFamilyByModelID(input);
if (!family) {
throw new Error(`No model ID or family found with ID: ${input}`);
}
return family;
}
return input;
};
isModelInFamily(modelID, familyID) {
const family = this.getFamilyByFamilyID(familyID);
return family?.modelIDs.includes(modelID) ?? false;
}
}