dsaf / ModelInfo.ts

Create ModelInfo.ts

1029ecd almost 5 years ago

6.19 kB

export class ModelInfo {
	/**
	 * Key to config.json file.
	 */
	key:          string;
	etag:         string;
	lastModified: Date;
	size:         number;
	modelId:      ModelId;
	author?:      AuthorId;
	siblings:     IS3ObjectWRelativeFilename[];
	config:       Obj;
	configTxt?:   string;  /// if flag is set when fetching.
	downloads?:   number;  /// if flag is set when fetching.
	naturalIdx:   number;
	cardSource?:  Source;
	cardData?:    Obj;

	constructor(o: Partial<ModelInfo>) {
		return Object.assign(this, o);
	}
	
	get jsonUrl(): string {
		return Bucket.R.models.urlForKey(this.key);
	}
	
	get cdnJsonUrl(): string {
		return Bucket.R.models.cdnUrlForKey(this.key);
	}
	
	async validate(): Promise<Ajv.ErrorObject[] | undefined> {
		const jsonSchema = JSON.parse(
			await fs.promises.readFile(CONFIG_JSON_SCHEMA, 'utf8')
		);
		const ajv = new Ajv();
		ajv.validate(jsonSchema, this.config);
		return ajv.errors ?? undefined;
	}
	
	/**
	 * Readme key, w. and w/o S3 prefix.
	 */
	get readmeKey(): string {
		return this.key.replace("config.json", "README.md");
	}
	get readmeTrimmedKey(): string {
		return Utils.trimPrefix(this.readmeKey, S3_MODELS_PREFIX);
	}
	
	/**
	 * ["pytorch", "tf", ...]
	 */
	get mlFrameworks(): string[] {
		return Object.keys(FileType).filter(k => {
			const filename = FileType[k];
			const isExtension = filename.startsWith(".");
			return isExtension 
				? this.siblings.some(sibling => sibling.rfilename.endsWith(filename))
				: this.siblings.some(sibling => sibling.rfilename === filename);
		});
	}
	/**
	 * What to display in the code sample.
	 */
	get autoArchitecture(): string {
		const useTF = this.mlFrameworks.includes("tf") && ! this.mlFrameworks.includes("pytorch");
		const arch = this.autoArchType[0];
		return useTF ? `TF${arch}` : arch;
	}
	get autoArchType(): [string, string | undefined] {
		const architectures = this.config.architectures;
		if (!architectures || architectures.length === 0) {
			return ["AutoModel", undefined];
		}
		const architecture = architectures[0].toString() as string;
		if (architecture.endsWith("ForQuestionAnswering")) {
			return ["AutoModelForQuestionAnswering", "question-answering"];
		}
		else if (architecture.endsWith("ForTokenClassification")) {
			return ["AutoModelForTokenClassification", "token-classification"];
		}
		else if (architecture.endsWith("ForSequenceClassification")) {
			return ["AutoModelForSequenceClassification", "text-classification"];
		}
		else if (architecture.endsWith("ForMultipleChoice")) {
			return ["AutoModelForMultipleChoice", "multiple-choice"];
		}
		else if (architecture.endsWith("ForPreTraining")) {
			return ["AutoModelForPreTraining", "pretraining"];
		}
		else if (architecture.endsWith("ForMaskedLM")) {
			return ["AutoModelForMaskedLM", "masked-lm"];
		}
		else if (architecture.endsWith("ForCausalLM")) {
			return ["AutoModelForCausalLM", "causal-lm"];
		}
		else if (
			   architecture.endsWith("ForConditionalGeneration")
			|| architecture.endsWith("MTModel")
			|| architecture == "EncoderDecoderModel"
		) {
			return ["AutoModelForSeq2SeqLM", "seq2seq"];
		}
		else if (architecture.includes("LMHead")) {
			return ["AutoModelWithLMHead", "lm-head"];
		}
		else if (architecture.endsWith("Model")) {
			return ["AutoModel", undefined];
		}
		else {
			return [architecture, undefined];
		}
	}
	/**
	 * All tags
	 */
	get tags(): string[] {
		const x = [
			...this.mlFrameworks,
		];
		if (this.config.model_type) {
			x.push(this.config.model_type);
		}
		const arch = this.autoArchType[1];
		if (arch) {
			x.push(arch);
		}
		if (arch === "lm-head" && this.config.model_type) {
			if ([
				"t5",
				"bart",
				"marian",
			].includes(this.config.model_type)) {
				x.push("seq2seq");
			}
			else if ([
				"gpt2",
				"ctrl",
				"openai-gpt",
				"xlnet",
				"transfo-xl",
				"reformer",
			].includes(this.config.model_type)) {
				x.push("causal-lm");
			}
			else {
				x.push("masked-lm");
			}
		}
		x.push(
			...this.languages() ?? []
		);
		x.push(
			...this.datasets().map(k => `dataset:${k}`)
		);
		for (let [k, v] of Object.entries(this.cardData ?? {})) {
			if (!['tags', 'license'].includes(k)) {
				/// ^^ whitelist of other accepted keys
				continue;
			}
			if (typeof v === 'string') {
				v = [ v ];
			} else if (Utils.isStrArray(v)) {
				/// ok
			} else {
				c.error(`Invalid ${k} tag type`, v);
				c.debug(this.modelId);
				continue;
			}
			if (k === 'license') {
				x.push(...v.map(x => `license:${x.toLowerCase()}`));
			} else {
				x.push(...v);
			}
		}
		if (this.config.task_specific_params) {
			const keys = Object.keys(this.config.task_specific_params);
			for (const key of keys) {
				x.push(`pipeline:${key}`);
			}
		}
		const explicit_ptag = this.cardData?.pipeline_tag;
		if (explicit_ptag) {
			if (typeof explicit_ptag === 'string') {
				x.push(`pipeline_tag:${explicit_ptag}`);
			} else {
				x.push(`pipeline_tag:invalid`);
			}
		}
		return [...new Set(x)];
	}
	
	get pipeline_tag(): (keyof typeof PipelineType) | undefined {
		if (isBlacklisted(this.modelId) || this.cardData?.inference === false) {
			return undefined;
		}
		
		const explicit_ptag = this.cardData?.pipeline_tag;
		if (explicit_ptag) {
			if (typeof explicit_ptag == 'string') {
				return explicit_ptag as keyof typeof PipelineType;
			} else {
				c.error(`Invalid explicit pipeline_tag`, explicit_ptag);
				return undefined;
			}
		}
		
		const tags = this.tags;
		/// Special case for translation
		/// Get the first of the explicit tags that matches.
		const EXPLICIT_PREFIX = "pipeline:";
		const explicit_tag = tags.find(x => x.startsWith(EXPLICIT_PREFIX + `translation`));
		if (!!explicit_tag) {
			return "translation";
		}
		/// Otherwise, get the first (most specific) match **from the mapping**.
		for (const ptag of ALL_PIPELINE_TYPES) {
			if (tags.includes(ptag)) {
				return ptag;
			}
		}
		/// Extra mapping
		const mapping = new Map<string, keyof typeof PipelineType>([
			["seq2seq", "text-generation"],
			["causal-lm", "text-generation"],
			["masked-lm", "fill-mask"],
		]);
		for (const [tag, ptag] of mapping) {
			if (tags.includes(tag)) {
				return ptag;
			}
		}
	}
}

	export class ModelInfo {
	/**
	* Key to config.json file.
	*/
	key: string;
	etag: string;
	lastModified: Date;
	size: number;
	modelId: ModelId;
	author?: AuthorId;
	siblings: IS3ObjectWRelativeFilename[];
	config: Obj;
	configTxt?: string; /// if flag is set when fetching.
	downloads?: number; /// if flag is set when fetching.
	naturalIdx: number;
	cardSource?: Source;
	cardData?: Obj;

	constructor(o: Partial<ModelInfo>) {
	return Object.assign(this, o);
	}

	get jsonUrl(): string {
	return Bucket.R.models.urlForKey(this.key);
	}

	get cdnJsonUrl(): string {
	return Bucket.R.models.cdnUrlForKey(this.key);
	}

	async validate(): Promise<Ajv.ErrorObject[] \| undefined> {
	const jsonSchema = JSON.parse(
	await fs.promises.readFile(CONFIG_JSON_SCHEMA, 'utf8')
	);
	const ajv = new Ajv();
	ajv.validate(jsonSchema, this.config);
	return ajv.errors ?? undefined;
	}

	/**
	* Readme key, w. and w/o S3 prefix.
	*/
	get readmeKey(): string {
	return this.key.replace("config.json", "README.md");
	}
	get readmeTrimmedKey(): string {
	return Utils.trimPrefix(this.readmeKey, S3_MODELS_PREFIX);
	}

	/**
	* ["pytorch", "tf", ...]
	*/
	get mlFrameworks(): string[] {
	return Object.keys(FileType).filter(k => {
	const filename = FileType[k];
	const isExtension = filename.startsWith(".");
	return isExtension
	? this.siblings.some(sibling => sibling.rfilename.endsWith(filename))
	: this.siblings.some(sibling => sibling.rfilename === filename);
	});
	}
	/**
	* What to display in the code sample.
	*/
	get autoArchitecture(): string {
	const useTF = this.mlFrameworks.includes("tf") && ! this.mlFrameworks.includes("pytorch");
	const arch = this.autoArchType[0];
	return useTF ? `TF${arch}` : arch;
	}
	get autoArchType(): [string, string \| undefined] {
	const architectures = this.config.architectures;
	if (!architectures \|\| architectures.length === 0) {
	return ["AutoModel", undefined];
	}
	const architecture = architectures[0].toString() as string;
	if (architecture.endsWith("ForQuestionAnswering")) {
	return ["AutoModelForQuestionAnswering", "question-answering"];
	}
	else if (architecture.endsWith("ForTokenClassification")) {
	return ["AutoModelForTokenClassification", "token-classification"];
	}
	else if (architecture.endsWith("ForSequenceClassification")) {
	return ["AutoModelForSequenceClassification", "text-classification"];
	}
	else if (architecture.endsWith("ForMultipleChoice")) {
	return ["AutoModelForMultipleChoice", "multiple-choice"];
	}
	else if (architecture.endsWith("ForPreTraining")) {
	return ["AutoModelForPreTraining", "pretraining"];
	}
	else if (architecture.endsWith("ForMaskedLM")) {
	return ["AutoModelForMaskedLM", "masked-lm"];
	}
	else if (architecture.endsWith("ForCausalLM")) {
	return ["AutoModelForCausalLM", "causal-lm"];
	}
	else if (
	architecture.endsWith("ForConditionalGeneration")
	\|\| architecture.endsWith("MTModel")
	\|\| architecture == "EncoderDecoderModel"
	) {
	return ["AutoModelForSeq2SeqLM", "seq2seq"];
	}
	else if (architecture.includes("LMHead")) {
	return ["AutoModelWithLMHead", "lm-head"];
	}
	else if (architecture.endsWith("Model")) {
	return ["AutoModel", undefined];
	}
	else {
	return [architecture, undefined];
	}
	}
	/**
	* All tags
	*/
	get tags(): string[] {
	const x = [
	...this.mlFrameworks,
	];
	if (this.config.model_type) {
	x.push(this.config.model_type);
	}
	const arch = this.autoArchType[1];
	if (arch) {
	x.push(arch);
	}
	if (arch === "lm-head" && this.config.model_type) {
	if ([
	"t5",
	"bart",
	"marian",
	].includes(this.config.model_type)) {
	x.push("seq2seq");
	}
	else if ([
	"gpt2",
	"ctrl",
	"openai-gpt",
	"xlnet",
	"transfo-xl",
	"reformer",
	].includes(this.config.model_type)) {
	x.push("causal-lm");
	}
	else {
	x.push("masked-lm");
	}
	}
	x.push(
	...this.languages() ?? []
	);
	x.push(
	...this.datasets().map(k => `dataset:${k}`)
	);
	for (let [k, v] of Object.entries(this.cardData ?? {})) {
	if (!['tags', 'license'].includes(k)) {
	/// ^^ whitelist of other accepted keys
	continue;
	}
	if (typeof v === 'string') {
	v = [ v ];
	} else if (Utils.isStrArray(v)) {
	/// ok
	} else {
	c.error(`Invalid ${k} tag type`, v);
	c.debug(this.modelId);
	continue;
	}
	if (k === 'license') {
	x.push(...v.map(x => `license:${x.toLowerCase()}`));
	} else {
	x.push(...v);
	}
	}
	if (this.config.task_specific_params) {
	const keys = Object.keys(this.config.task_specific_params);
	for (const key of keys) {
	x.push(`pipeline:${key}`);
	}
	}
	const explicit_ptag = this.cardData?.pipeline_tag;
	if (explicit_ptag) {
	if (typeof explicit_ptag === 'string') {
	x.push(`pipeline_tag:${explicit_ptag}`);
	} else {
	x.push(`pipeline_tag:invalid`);
	}
	}
	return [...new Set(x)];
	}

	get pipeline_tag(): (keyof typeof PipelineType) \| undefined {
	if (isBlacklisted(this.modelId) \|\| this.cardData?.inference === false) {
	return undefined;
	}

	const explicit_ptag = this.cardData?.pipeline_tag;
	if (explicit_ptag) {
	if (typeof explicit_ptag == 'string') {
	return explicit_ptag as keyof typeof PipelineType;
	} else {
	c.error(`Invalid explicit pipeline_tag`, explicit_ptag);
	return undefined;
	}
	}

	const tags = this.tags;
	/// Special case for translation
	/// Get the first of the explicit tags that matches.
	const EXPLICIT_PREFIX = "pipeline:";
	const explicit_tag = tags.find(x => x.startsWith(EXPLICIT_PREFIX + `translation`));
	if (!!explicit_tag) {
	return "translation";
	}
	/// Otherwise, get the first (most specific) match from the mapping.
	for (const ptag of ALL_PIPELINE_TYPES) {
	if (tags.includes(ptag)) {
	return ptag;
	}
	}
	/// Extra mapping
	const mapping = new Map<string, keyof typeof PipelineType>([
	["seq2seq", "text-generation"],
	["causal-lm", "text-generation"],
	["masked-lm", "fill-mask"],
	]);
	for (const [tag, ptag] of mapping) {
	if (tags.includes(tag)) {
	return ptag;
	}
	}
	}
	}