feat(models): adding backend logic to collect model metadata (#1352)

* feat(models): adding backend logic to collect model metadata Signed-off-by: axel7083 <[email protected]> * Apply suggestions from code review Signed-off-by: axel7083 <[email protected]> * fix: yarn.lock after rebase Signed-off-by: axel7083 <[email protected]> * fix: yarn.lock after rebase 2 Signed-off-by: axel7083 <[email protected]> * fix: yarn.lock after rebase 3 Signed-off-by: axel7083 <[email protected]> --------- Signed-off-by: axel7083 <[email protected]>
containers · Jul 10, 2024 · b78b9e6 · b78b9e6
1 parent 18ed74e
commit b78b9e6
Show file tree

Hide file tree

Showing 7 changed files with 154 additions and 1 deletion.
diff --git a/packages/backend/package.json b/packages/backend/package.json
@@ -55,6 +55,7 @@
     "watch": "vite --mode development build -w"
   },
   "dependencies": {
+    "@huggingface/gguf": "^0.1.7",
     "isomorphic-git": "^1.27.1",
     "mustache": "^4.2.0",
     "openai": "^4.52.3",

diff --git a/packages/backend/src/managers/modelsManager.spec.ts b/packages/backend/src/managers/modelsManager.spec.ts
@@ -29,6 +29,8 @@ import * as utils from '../utils/utils';
 import { TaskRegistry } from '../registries/TaskRegistry';
 import type { CancellationTokenRegistry } from '../registries/CancellationTokenRegistry';
 import * as sha from '../utils/sha';
+import type { GGUFParseOutput } from '@huggingface/gguf';
+import { gguf } from '@huggingface/gguf';
 
 const mocks = vi.hoisted(() => {
   return {
@@ -45,6 +47,10 @@ const mocks = vi.hoisted(() => {
   };
 });
 
+vi.mock('@huggingface/gguf', () => ({
+  gguf: vi.fn(),
+}));
+
 vi.mock('../utils/podman', () => ({
   getFirstRunningMachineName: mocks.getFirstRunningMachineNameMock,
   getPodmanCli: mocks.getPodmanCliMock,
@@ -833,3 +839,99 @@ describe('downloadModel', () => {
     expect(mocks.onEventDownloadMock).toHaveBeenCalledTimes(2);
   });
 });
+
+describe('getModelMetadata', () => {
+  test('unknown model', async () => {
+    const manager = new ModelsManager(
+      'appdir',
+      {} as Webview,
+      {
+        getModels: (): ModelInfo[] => [],
+      } as CatalogManager,
+      telemetryLogger,
+      taskRegistry,
+      cancellationTokenRegistryMock,
+    );
+
+    await expect(() => manager.getModelMetadata('unknown-model-id')).rejects.toThrowError(
+      'model with id unknown-model-id does not exists.',
+    );
+  });
+
+  test('remote model', async () => {
+    const manager = new ModelsManager(
+      'appdir',
+      {} as Webview,
+      {
+        getModels: (): ModelInfo[] => [
+          {
+            id: 'test-model-id',
+            url: 'dummy-url',
+            file: undefined,
+          } as unknown as ModelInfo,
+        ],
+        onCatalogUpdate: vi.fn(),
+      } as unknown as CatalogManager,
+      telemetryLogger,
+      taskRegistry,
+      cancellationTokenRegistryMock,
+    );
+
+    manager.init();
+
+    const fakeMetadata: Record<string, string> = {
+      hello: 'world',
+    };
+
+    vi.mocked(gguf).mockResolvedValue({
+      metadata: fakeMetadata,
+    } as unknown as GGUFParseOutput & { parameterCount: number });
+
+    const result = await manager.getModelMetadata('test-model-id');
+    expect(result).toStrictEqual(fakeMetadata);
+
+    expect(gguf).toHaveBeenCalledWith('dummy-url');
+  });
+
+  test('local model', async () => {
+    const manager = new ModelsManager(
+      'appdir',
+      {
+        postMessage: vi.fn(),
+      } as unknown as Webview,
+      {
+        getModels: (): ModelInfo[] => [
+          {
+            id: 'test-model-id',
+            url: 'dummy-url',
+            file: {
+              file: 'random',
+              path: 'dummy-path',
+            },
+          } as unknown as ModelInfo,
+        ],
+        onCatalogUpdate: vi.fn(),
+      } as unknown as CatalogManager,
+      telemetryLogger,
+      taskRegistry,
+      cancellationTokenRegistryMock,
+    );
+
+    manager.init();
+
+    const fakeMetadata: Record<string, string> = {
+      hello: 'world',
+    };
+
+    vi.mocked(gguf).mockResolvedValue({
+      metadata: fakeMetadata,
+    } as unknown as GGUFParseOutput & { parameterCount: number });
+
+    const result = await manager.getModelMetadata('test-model-id');
+    expect(result).toStrictEqual(fakeMetadata);
+
+    expect(gguf).toHaveBeenCalledWith(path.join('dummy-path', 'random'), {
+      allowLocalFile: true,
+    });
+  });
+});
diff --git a/packages/backend/src/managers/modelsManager.ts b/packages/backend/src/managers/modelsManager.ts
@@ -17,7 +17,7 @@
  ***********************************************************************/
 
 import type { LocalModelInfo } from '@shared/src/models/ILocalModelInfo';
-import fs from 'fs';
+import fs from 'node:fs';
 import * as path from 'node:path';
 import { type Webview, fs as apiFs, type Disposable, env } from '@podman-desktop/api';
 import { Messages } from '@shared/Messages';
@@ -34,6 +34,8 @@ import { deleteRemoteModel, getLocalModelFile, isModelUploaded } from '../utils/
 import { getFirstRunningMachineName } from '../utils/podman';
 import type { CancellationTokenRegistry } from '../registries/CancellationTokenRegistry';
 import { hasValidSha } from '../utils/sha';
+import type { GGUFParseOutput } from '@huggingface/gguf';
+import { gguf } from '@huggingface/gguf';
 
 export class ModelsManager implements Disposable {
   #models: Map<string, ModelInfo>;
@@ -422,4 +424,34 @@ export class ModelsManager implements Disposable {
     // perform download
     return uploader.perform(model.id);
   }
+
+  async getModelMetadata(modelId: string): Promise<Record<string, unknown>> {
+    const model = this.#models.get(modelId);
+    if (!model) throw new Error(`model with id ${modelId} does not exists.`);
+
+    const before = performance.now();
+    const data: Record<string, unknown> = {
+      'model-id': model.url ? modelId : 'imported', // filter imported models
+    };
+
+    try {
+      let result: GGUFParseOutput<{ strict: false }>;
+      if (this.isModelOnDisk(modelId)) {
+        const modelPath = path.normalize(getLocalModelFile(model));
+        result = await gguf(modelPath, { allowLocalFile: true });
+      } else if (model.url) {
+        result = await gguf(model.url);
+      } else {
+        throw new Error('cannot get model metadata');
+      }
+      return result.metadata;
+    } catch (err: unknown) {
+      data['error'] = err;
+      console.error(err);
+      throw err;
+    } finally {
+      data['duration'] = performance.now() - before;
+      this.telemetry.logUsage('get-metadata', data);
+    }
+  }
 }
diff --git a/packages/backend/src/studio-api-impl.ts b/packages/backend/src/studio-api-impl.ts
@@ -221,6 +221,10 @@ export class StudioApiImpl implements StudioAPI {
     return this.modelsManager.getModelsInfo();
   }
 
+  getModelMetadata(modelId: string): Promise<Record<string, unknown>> {
+    return this.modelsManager.getModelMetadata(modelId);
+  }
+
   async getCatalog(): Promise<ApplicationCatalog> {
     return this.catalogManager.getCatalog();
   }

diff --git a/packages/backend/vite.config.js b/packages/backend/vite.config.js
@@ -49,6 +49,7 @@ const config = {
     rollupOptions: {
       external: [
         '@podman-desktop/api',
+        '@huggingface/gguf',
         ...builtinModules.flatMap(p => [p, `node:${p}`]),
       ],
       output: {

diff --git a/packages/shared/src/StudioAPI.ts b/packages/shared/src/StudioAPI.ts
@@ -58,6 +58,14 @@ export abstract class StudioAPI {
    * Get the information of models saved locally into the user's directory
    */
   abstract getModelsInfo(): Promise<ModelInfo[]>;
+
+  /**
+   * Given a modelId will return the model metadata
+   * @remark If the model is not available locally, a fetch request will be used to get its metadata from the header.
+   * @param modelId
+   */
+  abstract getModelMetadata(modelId: string): Promise<Record<string, unknown>>;
+
   /**
    * Delete the folder containing the model from local storage
    */

diff --git a/yarn.lock b/yarn.lock
@@ -281,6 +281,11 @@
   dependencies:
     "@fortawesome/fontawesome-common-types" "6.5.2"
 
+"@huggingface/gguf@^0.1.7":
+  version "0.1.7"
+  resolved "https://registry.yarnpkg.com/@huggingface/gguf/-/gguf-0.1.7.tgz#0f5afab60f182cf27341738bff794750cacb5876"
+  integrity sha512-RQN1WwuusLjiBTNFuAJCUlhRejIhKt395ywnTmc+Jy8dajGwk8k7EsfgtxVqkBSTWy9D55XzCII7jUjkHEv3JA==
+
 "@humanwhocodes/config-array@^0.11.14":
   version "0.11.14"
   resolved "https://registry.yarnpkg.com/@humanwhocodes/config-array/-/config-array-0.11.14.tgz#d78e481a039f7566ecc9660b4ea7fe6b1fec442b"