refactor: move node specific code into separate file

microsoft · Apr 15, 2024 · 1be1a2b · 1be1a2b
1 parent bb49b83
commit 1be1a2b
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 36 deletions.
diff --git a/lib/index.node.ts b/lib/index.node.ts
@@ -0,0 +1,48 @@
+import { ModelOperations, ModelOperationsOptions } from './index';
+export * from './index'
+
+type NodeModelOperationsOptions = Omit<ModelOperationsOptions, 'modelJsonLoaderFunc' | 'weightsLoaderFunc'> & Partial<ModelOperationsOptions>
+
+class NodeModelOperations extends ModelOperations {
+	private static NODE_MODEL_JSON_FUNC: () => Promise<{ [key:string]: any }> = async () => {
+		const fs = await import('fs');
+		const path = await import('path');
+
+		return new Promise<any>((resolve, reject) => {
+			fs.readFile(path.join(__dirname, '..', '..', 'model', 'model.json'), (err, data) => {
+				if(err) {
+					reject(err);
+					return;
+				}
+				resolve(JSON.parse(data.toString()));
+			});
+		});
+	}
+
+	private static NODE_WEIGHTS_FUNC: () => Promise<ArrayBuffer> = async () => {
+		const fs = await import('fs');
+		const path = await import('path');
+
+		return new Promise<ArrayBuffer>((resolve, reject) => {
+			fs.readFile(path.join(__dirname, '..', '..', 'model', 'group1-shard1of1.bin'), (err, data) => {
+				if(err) {
+					reject(err);
+					return;
+				}
+				resolve(data.buffer);
+			});
+		});
+	}
+
+	constructor(modelOptions?: NodeModelOperationsOptions) {
+		super({
+			modelJsonLoaderFunc: modelOptions?.modelJsonLoaderFunc ?? NodeModelOperations.NODE_MODEL_JSON_FUNC,
+			weightsLoaderFunc: modelOptions?.weightsLoaderFunc ?? NodeModelOperations.NODE_WEIGHTS_FUNC,
+			...modelOptions
+		})
+	}
+}
+
+export {
+	NodeModelOperations as ModelOperations
+}
diff --git a/lib/index.ts b/lib/index.ts
@@ -71,8 +71,8 @@ class InMemoryIOHandler implements io.IOHandler {
 }
 
 export interface ModelOperationsOptions {
-	modelJsonLoaderFunc?: () => Promise<{ [key:string]: any }>;
-	weightsLoaderFunc?: () => Promise<ArrayBuffer>;
+	modelJsonLoaderFunc: () => Promise<{ [key:string]: any }>;
+	weightsLoaderFunc: () => Promise<ArrayBuffer>;
 	minContentSize?: number;
 	maxContentSize?: number;
 	normalizeNewline?: boolean;
@@ -82,36 +82,6 @@ export class ModelOperations {
 	private static DEFAULT_MAX_CONTENT_SIZE = 100000;
 	private static DEFAULT_MIN_CONTENT_SIZE = 20;
 
-	private static NODE_MODEL_JSON_FUNC: () => Promise<{ [key:string]: any }> = async () => {
-		const fs = await import('fs');
-		const path = await import('path');
-
-		return new Promise<any>((resolve, reject) => {
-			fs.readFile(path.join(__dirname, '..', '..', 'model', 'model.json'), (err, data) => {
-				if(err) {
-					reject(err);
-					return;
-				}
-				resolve(JSON.parse(data.toString()));
-			});
-		});
-	}
-
-	private static NODE_WEIGHTS_FUNC: () => Promise<ArrayBuffer> = async () => {
-		const fs = await import('fs');
-		const path = await import('path');
-
-		return new Promise<ArrayBuffer>((resolve, reject) => {
-			fs.readFile(path.join(__dirname, '..', '..', 'model', 'group1-shard1of1.bin'), (err, data) => {
-				if(err) {
-					reject(err);
-					return;
-				}
-				resolve(data.buffer);
-			});
-		});
-	}
-
 	private _model: GraphModel | undefined;
 	private _modelJson: io.ModelJSON | undefined;
 	private _weights: ArrayBuffer | undefined;
@@ -121,9 +91,9 @@ export class ModelOperations {
 	private readonly _weightsLoaderFunc: () => Promise<ArrayBuffer>;
 	private readonly _normalizeNewline: boolean;
 
-	constructor(modelOptions?: ModelOperationsOptions) {
-		this._modelJsonLoaderFunc = modelOptions?.modelJsonLoaderFunc ?? ModelOperations.NODE_MODEL_JSON_FUNC;
-		this._weightsLoaderFunc = modelOptions?.weightsLoaderFunc ?? ModelOperations.NODE_WEIGHTS_FUNC;
+	constructor(modelOptions: ModelOperationsOptions) {
+		this._modelJsonLoaderFunc = modelOptions?.modelJsonLoaderFunc;
+		this._weightsLoaderFunc = modelOptions?.weightsLoaderFunc;
 		this._minContentSize = modelOptions?.minContentSize ?? ModelOperations.DEFAULT_MIN_CONTENT_SIZE;
 		this._maxContentSize = modelOptions?.maxContentSize ?? ModelOperations.DEFAULT_MAX_CONTENT_SIZE;
 		this._normalizeNewline = modelOptions?.normalizeNewline ?? true;

diff --git a/rollup.config.ts b/rollup.config.ts
@@ -14,7 +14,7 @@ const isProduction = process.env.NODE_ENV === 'production'
 
 export default rollup.defineConfig([{
   input: {
-    index: './lib/index.ts'
+    index: './lib/index.node.ts'
   },
   output: [{
     sourcemap: true,