fix: error handling for model imports should be handled gracefully

janhq · Oct 4, 2024 · 2c2f179 · 2c2f179
1 parent 27c6ac6
commit 2c2f179
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 8 deletions.
diff --git a/extensions/model-extension/src/index.test.ts b/extensions/model-extension/src/index.test.ts
@@ -784,5 +784,63 @@ describe('JanModelExtension', () => {
         expect.anything()
       )
     })
+
+    it('should handle model with valid chat_template', async () => {
+      executeMock.mockResolvedValue('{prompt}')
+      ;(gguf as jest.Mock).mockResolvedValue({
+        metadata: {},
+      })
+      // @ts-ignore
+      global.NODE = 'node'
+      // @ts-ignore
+      global.DEFAULT_MODEL = {
+        parameters: { stop: [] },
+        settings: {
+          prompt_template: '<|im-start|>{prompt}<|im-end|>',
+        },
+      }
+
+      const result = await sut.retrieveGGUFMetadata({})
+
+      expect(result).toEqual({
+        parameters: {
+          stop: [],
+        },
+        settings: {
+          ctx_len: 4096,
+          ngl: 33,
+          prompt_template: '{prompt}',
+        },
+      })
+    })
+
+    it('should handle model without chat_template', async () => {
+      executeMock.mockRejectedValue({})
+      ;(gguf as jest.Mock).mockResolvedValue({
+        metadata: {},
+      })
+      // @ts-ignore
+      global.NODE = 'node'
+      // @ts-ignore
+      global.DEFAULT_MODEL = {
+        parameters: { stop: [] },
+        settings: {
+          prompt_template: '<|im-start|>{prompt}<|im-end|>',
+        },
+      }
+
+      const result = await sut.retrieveGGUFMetadata({})
+
+      expect(result).toEqual({
+        parameters: {
+          stop: [],
+        },
+        settings: {
+          ctx_len: 4096,
+          ngl: 33,
+          prompt_template: '<|im-start|>{prompt}<|im-end|>',
+        },
+      })
+    })
   })
 })
diff --git a/extensions/model-extension/src/index.ts b/extensions/model-extension/src/index.ts
@@ -577,7 +577,7 @@ export default class JanModelExtension extends ModelExtension {
         dirName,
         binaryFileName,
       ])
-    )
+    ).catch(() => undefined)
 
     const updatedModel = await this.retrieveGGUFMetadata(metadata)
 
@@ -877,8 +877,13 @@ export default class JanModelExtension extends ModelExtension {
    * @returns
    */
   async retrieveGGUFMetadata(metadata: any): Promise<Partial<Model>> {
-    const template = await executeOnMain(NODE, 'renderJinjaTemplate', metadata)
     const defaultModel = DEFAULT_MODEL as Model
+    var template = await executeOnMain(
+      NODE,
+      'renderJinjaTemplate',
+      metadata
+    ).catch(() => undefined)
+
     const eos_id = metadata['tokenizer.ggml.eos_token_id']
     const architecture = metadata['general.architecture']
 

diff --git a/extensions/model-extension/src/node/index.ts b/extensions/model-extension/src/node/index.ts
@@ -16,12 +16,7 @@ export const retrieveGGUFMetadata = async (ggufPath: string) => {
     // Parse metadata and tensor info
     const { metadata } = ggufMetadata(buffer.buffer)
 
-    // Parse jinja template
-    const renderedTemplate = renderJinjaTemplate(metadata)
-    return {
-      ...metadata,
-      parsed_chat_template: renderedTemplate,
-    }
+    return metadata
   } catch (e) {
     console.log('[MODEL_EXT]', e)
   }