fixed loadweightmap always fetching from relative url

justadudewhohacks · justadudewhohacks · commit 656602f746a4 · 2018-07-26T10:27:41.000+02:00
diff --git a/src/commons/loadWeightMap.ts b/src/commons/loadWeightMap.ts
@@ -1,24 +1,35 @@
 import * as tf from '@tensorflow/tfjs-core';
 
 export function getModelUris(uri: string | undefined, defaultModelName: string) {
-  const parts = (uri || '').split('/')
-
-  const modelBaseUri = (
-    (uri || '').endsWith('.json')
-      ? parts.slice(0, parts.length - 1)
-      : parts
-  ).filter(s => s).join('/')
-
   const defaultManifestFilename = `${defaultModelName}-weights_manifest.json`
-  const manifestUri = !uri || !modelBaseUri
-    ? defaultManifestFilename
-    : (
-      uri.endsWith('.json')
-        ? uri
-        : `${modelBaseUri}/${defaultManifestFilename}`
-    )
-
-  return { manifestUri, modelBaseUri }
+
+  if (!uri) {
+    return {
+      modelBaseUri: '',
+      manifestUri: defaultManifestFilename
+    }
+  }
+
+  if (uri === '/') {
+    return {
+      modelBaseUri: '/',
+      manifestUri: `/${defaultManifestFilename}`
+    }
+  }
+
+  const parts = uri.split('/').filter(s => s)
+
+  const manifestFile = uri.endsWith('.json')
+    ? parts[parts.length - 1]
+    : defaultManifestFilename
+
+  let modelBaseUri = (uri.endsWith('.json') ? parts.slice(0, parts.length - 1) : parts).join('/')
+  modelBaseUri = uri.startsWith('/') ? `/${modelBaseUri}` : modelBaseUri
+
+  return {
+    modelBaseUri,
+    manifestUri: modelBaseUri === '/' ? `/${manifestFile}` : `${modelBaseUri}/${manifestFile}`
+  }
 }
 
 export async function loadWeightMap(
diff --git a/test/tests/commons/loadWeightMap.test.ts b/test/tests/commons/loadWeightMap.test.ts
@@ -1,30 +1,30 @@
 import { getModelUris } from '../../../src/commons/loadWeightMap';
 
-const FAKE_DEFAULT_MODEL_NAME = 'default_model_name'
+const FAKE_DEFAULT_MODEL_NAME = 'fake_model_name'
 
 describe('loadWeightMap', () => {
 
   describe('getModelUris', () => {
 
-    it('returns uris from top level url if no argument passed', () => {
+    it('returns uris from relative url if no argument passed', () => {
       const result = getModelUris(undefined, FAKE_DEFAULT_MODEL_NAME)
 
       expect(result.manifestUri).toEqual(`${FAKE_DEFAULT_MODEL_NAME}-weights_manifest.json`)
       expect(result.modelBaseUri).toEqual('')
     })
 
-    it('returns uris from top level url for empty string', () => {
+    it('returns uris from relative url for empty string', () => {
       const result = getModelUris('', FAKE_DEFAULT_MODEL_NAME)
 
       expect(result.manifestUri).toEqual(`${FAKE_DEFAULT_MODEL_NAME}-weights_manifest.json`)
       expect(result.modelBaseUri).toEqual('')
     })
 
-    it('returns uris for top level url', () => {
+    it('returns uris for top level url, leading slash preserved', () => {
       const result = getModelUris('/', FAKE_DEFAULT_MODEL_NAME)
 
-      expect(result.manifestUri).toEqual(`${FAKE_DEFAULT_MODEL_NAME}-weights_manifest.json`)
-      expect(result.modelBaseUri).toEqual('')
+      expect(result.manifestUri).toEqual(`/${FAKE_DEFAULT_MODEL_NAME}-weights_manifest.json`)
+      expect(result.modelBaseUri).toEqual('/')
     })
 
     it('returns uris, given url path', () => {
@@ -35,8 +35,8 @@ describe('loadWeightMap', () => {
       expect(result.modelBaseUri).toEqual(uri)
     })
 
-    it('returns uris, given url path, leading slash', () => {
-      const uri = 'path/to/modelfiles'
+    it('returns uris, given url path, leading slash preserved', () => {
+      const uri = '/path/to/modelfiles'
       const result = getModelUris(`/${uri}`, FAKE_DEFAULT_MODEL_NAME)
 
       expect(result.manifestUri).toEqual(`${uri}/${FAKE_DEFAULT_MODEL_NAME}-weights_manifest.json`)
@@ -51,6 +51,14 @@ describe('loadWeightMap', () => {
       expect(result.modelBaseUri).toEqual('path/to/modelfiles')
     })
 
+    it('returns uris, given manifest uri, leading slash preserved', () => {
+      const uri = '/path/to/modelfiles/model-weights_manifest.json'
+      const result = getModelUris(uri, FAKE_DEFAULT_MODEL_NAME)
+
+      expect(result.manifestUri).toEqual(uri)
+      expect(result.modelBaseUri).toEqual('/path/to/modelfiles')
+    })
+
   })
 
 })
