XciD's picture
XciD HF staff
initial commit
8969f81
raw history blame
No virus
3.32 kB
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;