server/models/Doc.ts

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>): 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<Doc> {
		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 `<code>` 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(/</g, "&lt;") ?? "";
			if (op.attributes && op.attributes.bold) {
				return `<strong>${escaped}</strong>`;
			}
			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(/</g, "&lt;");
	}
	
	duck(): string {
		return '___quack___';
	}
}


export default Doc;