import { ObjectID } from 'mongodb'; import { MongoObject } from './MongoObject'; import { Utils } from '../lib/Utils'; import { config } from '../lib/config'; import db from './Database'; export enum ModelId { 'distil-gpt2' = 'distil-gpt2', 'arxiv-nlp' = 'arxiv-nlp', 'gpt2-large' = 'gpt2-large', 'gpt2-xl' = 'gpt2-xl', 'xlnet' = 'xlnet', 'gpt' = 'gpt', 'ctrl' = 'ctrl', 'pplm' = 'pplm', } export const ALL_MODELS = Object.values(ModelId) as string[]; namespace Doc { interface InsertDeltaOperation { insert: string; attributes?: { link?: string; bold?: true; }; } export interface Contents { ops: InsertDeltaOperation[]; } } class Doc extends MongoObject { protected static __type = Doc; protected static __collectionName: string = 'docs'; protected static __idField: string = 'shortId'; _id: ObjectID; shortId: string; longId?: string; model: ModelId; contents: Doc.Contents; /** * Optional reference to another doc's `shortId`. */ clonedFrom?: string; title?: string; /** * Whether to display it on the front page. */ featured?: boolean; static fromObject(o: Partial): Doc { return Object.assign(new Doc(), o); } /** * Initialize a new doc with random ids. */ static seed(model: ModelId): Doc { return this.fromObject({ shortId: Utils.randomStr(10), longId: Utils.randomStr(24), model }); } /** * Insert a new doc duplicated from an existing one. */ async duplicate(): Promise { const newdoc = Doc.seed(this.model); newdoc.contents = this.contents; newdoc.title = this.title; newdoc.clonedFrom = this.shortId; await db.docs.insertOne(newdoc); return newdoc; } /** * This is displayed next to the doc title in a `` tag. */ get modelFamily(): string { if ([ 'arxiv-nlp', 'gpt2-large', 'gpt2-xl', ].includes(this.model)) { return 'gpt2'; } if (this.model === 'ctrl') { return `salesforce/ctrl`; } if (this.model === 'pplm') { return `uber/pplm`; } return this.model; } get hasCustomLogo(): boolean { return ['ctrl', 'pplm'].includes(this.model); } get modelInfoLink(): string { if (this.model === 'pplm') { return `https://github.com/huggingface/transformers/tree/master/examples/pplm`; } return `https://github.com/huggingface/transformers`; } get shareUrl(): string { return `${config.transformerAutocompleteUrl}/share/${this.shortId}`; } get editUrl(): string { return `${config.transformerAutocompleteUrl}/doc/${this.model}/${this.longId}/edit`; } /** * Construct a (safe) HTML representation of the quill delta content format. */ get getHTML(): string | undefined { if (! this.contents) { return ; } return this.contents.ops.map(op => { const escaped = op.insert?.toString().replace(/${escaped}`; } return escaped; }).join(""); } /** * Construct a (safe) JSON representation of the doc.contents. */ get contentsJson(): string | undefined { if (! this.contents) { return ; } return JSON.stringify(this.contents).replace(/