Create ModelInfo.ts

ModelInfo.ts

@@ -0,0 +1,234 @@
+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;
+			}
+		}
+	}
+}