feat: pull model yaml from hf

Author: namchuai

cortex-js/package.json

@@ -26,6 +26,7 @@
   },
   "dependencies": {
     "@huggingface/gguf": "^0.1.5",
+    "@huggingface/hub": "^0.15.1",
     "@nestjs/axios": "^3.0.2",
     "@nestjs/common": "^10.0.0",
     "@nestjs/config": "^3.2.2",
@@ -47,7 +48,8 @@
     "sqlite": "^5.1.1",
     "sqlite3": "^5.1.7",
     "typeorm": "^0.3.20",
-    "ulid": "^2.3.0"
+    "ulid": "^2.3.0",
+    "yaml": "^2.4.2"
   },
   "devDependencies": {
     "@nestjs/cli": "^10.0.0",

cortex-js/src/infrastructure/commanders/models/model-pull.command.ts

@@ -1,25 +1,151 @@
 import { CommandRunner, SubCommand } from 'nest-commander';
-import { exit } from 'node:process';
+import { exit, stdin, stdout } from 'node:process';
 import { ModelsCliUsecases } from '../usecases/models.cli.usecases';
+import { RepoDesignation, listFiles } from '@huggingface/hub';
+import YAML from 'yaml';
+import * as readline from 'node:readline/promises';
+import { basename } from 'node:path';
 
 @SubCommand({
   name: 'pull',
   aliases: ['download'],
   description: 'Download a model. Working with HuggingFace model id.',
 })
 export class ModelPullCommand extends CommandRunner {
+  private metadataFileName = 'metadata.yaml';
+
   constructor(private readonly modelsCliUsecases: ModelsCliUsecases) {
     super();
   }
 
   async run(input: string[]) {
     if (input.length < 1) {
-      console.error('Model ID is required');
+      console.error('Model Id is required');
       exit(1);
     }
 
-    await this.modelsCliUsecases.pullModel(input[0]);
+    // Check if metadata.yaml file exist
+    const metadata = await this.getJanMetadata(input[0]);
+
+    if (!metadata) {
+      await this.modelsCliUsecases.pullModel(input[0]);
+    } else {
+      await this.handleJanHqModel(input[0], metadata);
+    }
+
     console.log('\nDownload complete!');
     exit(0);
   }
+
+  private async getJanMetadata(input: string): Promise<any> {
+    // try to append with janhq/ if it's not already
+    const sanitizedInput = input.trim().startsWith('janhq/')
+      ? input
+      : `janhq/${input}`;
+
+    const repo: RepoDesignation = { type: 'model', name: sanitizedInput };
+    let isMetadataFileExist = false;
+    for await (const fileInfo of listFiles({ repo })) {
+      if (fileInfo.path === this.metadataFileName) {
+        isMetadataFileExist = true;
+        break;
+      }
+    }
+
+    if (!isMetadataFileExist) {
+      return undefined;
+    }
+
+    const path = `https://huggingface.co/${sanitizedInput}/raw/main/${this.metadataFileName}`;
+    const res = await fetch(path);
+    const metadataJson = await res.text();
+    const parsedMetadata = YAML.parse(metadataJson);
+    return parsedMetadata;
+  }
+
+  private async versionInquiry(tags: string[]): Promise<string> {
+    return new Promise((resolve) => {
+      let selectedTag = 'default';
+      let prompt = 'Select the version you want to download:\n';
+      for (let i = 0; i < tags.length; i++) {
+        prompt += `${i}. ${tags[i]}\n`;
+      }
+      prompt += '>> ';
+
+      const rl = readline.createInterface({
+        input: stdin,
+        output: stdout,
+        prompt: prompt,
+      });
+      rl.prompt();
+
+      rl.on('close', () => {
+        resolve(selectedTag);
+      });
+
+      rl.on('line', (input) => {
+        if (input.trim().length === 0) {
+          rl.close();
+        }
+
+        try {
+          if (Number(input) >= 0 && Number(input) < tags.length) {
+            selectedTag = tags[Number(input)];
+            rl.close();
+          } else {
+            console.error('Invalid version');
+            rl.prompt();
+          }
+        } catch (e) {
+          console.error('Invalid version');
+          rl.prompt();
+        }
+      });
+    });
+  }
+
+  private async handleJanHqModel(repoName: string, metadata: any) {
+    const sanitizedRepoName = repoName.trim().startsWith('janhq/')
+      ? repoName
+      : `janhq/${repoName}`;
+
+    const tags = metadata.tags;
+    let selectedTag = 'default';
+    const allTags: string[] = Object.keys(tags);
+
+    if (allTags.length > 1) {
+      selectedTag = await this.versionInquiry(allTags);
+    }
+
+    const branch = selectedTag;
+    const engine = 'llamacpp'; // TODO: currently, we only support llamacpp
+
+    const revision = metadata.tags?.[branch]?.[engine];
+    if (!revision) {
+      console.error("Can't find model revision.");
+      exit(1);
+    }
+
+    const repo: RepoDesignation = { type: 'model', name: sanitizedRepoName };
+    let ggufUrl: string | undefined = undefined;
+    for await (const fileInfo of listFiles({
+      repo: repo,
+      revision: revision,
+    })) {
+      if (fileInfo.path.endsWith('.gguf')) {
+        ggufUrl = `https://huggingface.co/${sanitizedRepoName}/resolve/${revision}/${fileInfo.path}`;
+        break;
+      }
+    }
+
+    if (!ggufUrl) {
+      console.error("Can't find model file.");
+      exit(1);
+    }
+    console.log('Downloading', basename(ggufUrl));
+    await this.modelsCliUsecases.pullModelWithExactUrl(
+      `${sanitizedRepoName}/${revision}`,
+      ggufUrl,
+    );
+  }
 }

cortex-js/src/infrastructure/commanders/usecases/models.cli.usecases.ts

@@ -98,6 +98,41 @@ export class ModelsCliUsecases {
     return this.modelsUsecases.remove(modelId);
   }
 
+  async pullModelWithExactUrl(modelId: string, url: string) {
+    const model: CreateModelDto = {
+      sources: [
+        {
+          url: url,
+        },
+      ],
+      id: modelId,
+      name: modelId,
+      version: '', // TODO: get version from the file
+      format: ModelFormat.GGUF,
+      description: '',
+      settings: {
+        prompt_template: '', // TODO: get prompt template from the file
+      },
+      parameters: {
+        stop: [],
+      },
+      metadata: {
+        author: 'janhq', // TODO: get author from the file
+        size: 0, // TODO: get size from the file
+        tags: [],
+      },
+      engine: 'cortex',
+    };
+    if (!(await this.modelsUsecases.findOne(modelId)))
+      await this.modelsUsecases.create(model);
+    const bar = new SingleBar({}, Presets.shades_classic);
+    bar.start(100, 0);
+    const callback = (progress: number) => {
+      bar.update(progress);
+    };
+    await this.modelsUsecases.downloadModel(modelId, callback);
+  }
+
   async pullModel(modelId: string) {
     if (modelId.includes('/')) {
       await this.pullHuggingFaceModel(modelId);