feat: implement ai model download using web workers

Author: joelshejar

lefthook.yml

@@ -5,7 +5,7 @@ pre-commit:
       run: pnpm biome check --write --unsafe --staged --no-errors-on-unmatched && git add -u
     typecheck:
       run: pnpm tsc
-    build:
-      run: pnpm build
+    # build:
+    #   run: pnpm build
     test:
       run: pnpm test:ci

package.json

@@ -1,7 +1,7 @@
 {
-	"name": "@codethicket/react-ai-translator",
+	"name": "@joelshejar/react-ai-translator",
 	"description": "",
-	"version": "0.0.0",
+	"version": "0.0.57",
 	"author": "Joel Rajesh",
 	"license": "MIT",
 	"keywords": [],
@@ -23,10 +23,10 @@
 		"link:self": "pnpm link --global",
 		"prepare": "lefthook install"
 	},
+	"module": "./dist/index.mjs",
 	"types": "./dist/index.d.ts",
 	"exports": {
 		".": {
-			"require": "./dist/index.js",
 			"import": "./dist/index.mjs"
 		}
 	},

src/hooks/index.ts renamed to src/hooks/index.js

@@ -0,0 +1,116 @@
+import { useEffect, useRef, useState } from "react";
+
+const useTranslation = (_workerScript) => {
+	const worker = useRef(null);
+
+	// States for managing translation
+	const [loading, setLoading] = useState(false);
+	const [modelLoading, setModelLoading] = useState(false);
+	const [progress, setProgress] = useState([]);
+	const [translatedText, setTranslatedText] = useState("");
+	const [error, setError] = useState(null);
+
+	console.log("hhhhh");
+
+	useEffect(() => {
+		if (!worker.current) {
+			// Initialize the worker
+			worker.current = new Worker(new URL("./worker.mjs", import.meta.url), {
+				type: "module",
+			});
+		}
+
+		// Handle worker messages
+		const onMessage = (e) => {
+			const { status, output, progress: progressValue, file } = e.data;
+
+			console.log(status, "status");
+
+			switch (status) {
+				case "initiate": {
+					console.log("initiate");
+					setLoading(true);
+					setModelLoading(true);
+					setProgress((prev) => [...prev, { file, progress: progressValue }]);
+					break;
+				}
+
+				case "progress": {
+					console.log("progress");
+					setProgress((prev) =>
+						prev.map((item) =>
+							item.file === file ? { ...item, progress: progressValue } : item,
+						),
+					);
+					break;
+				}
+
+				case "done": {
+					console.log("done");
+					setModelLoading(false);
+					setProgress((prev) => prev.filter((item) => item.file !== file));
+					break;
+				}
+
+				case "ready": {
+					console.log("ready");
+					setLoading(false);
+					break;
+				}
+
+				case "update": {
+					console.log("update");
+					console.log(output, "output");
+					// Append partial translations
+					setTranslatedText(output);
+					break;
+				}
+
+				case "complete": {
+					console.log("complete");
+					setLoading(false);
+					break;
+				}
+			}
+		};
+
+		worker.current.addEventListener("message", onMessage);
+
+		return () => {
+			// worker.current.terminate(); // Clean up worker on unmount
+			// worker.current = null;
+		};
+	}, []);
+
+	// Function to send translation requests to the worker
+	const translate = (text, srcLang, tgtLang) => {
+		console.log("kkkk");
+		if (!worker.current) {
+			console.error("Worker is not initialized.");
+			return;
+		}
+
+		setLoading(true);
+		setError(null);
+		setTranslatedText(""); // Reset the output
+
+		console.log("lllll", text, srcLang, tgtLang);
+
+		worker.current.postMessage({
+			text,
+			src_lang: srcLang,
+			tgt_lang: tgtLang,
+		});
+	};
+
+	return {
+		translate,
+		translatedText,
+		loading,
+		modelLoading,
+		progress,
+		error,
+	};
+};
+
+export default useTranslation;

src/hooks/useTranslation.js

@@ -1,6 +1,6 @@
 import type { Meta, StoryFn } from "@storybook/react";
 
-import { Example } from "..";
+import { Example } from "../Example";
 
 export default {
 	title: "Example",

src/hooks/useTranslation.ts

@@ -5,8 +5,6 @@ import { pipeline } from "@xenova/transformers";
  * pipeline is loaded. This is because loading the pipeline is an expensive
  * operation and we don't want to do it every time we want to translate a sentence.
  */
-
-// biome-disable-next-line complexity/noStaticOnlyClass
 class MyTranslationPipeline {
 	static task = "translation";
 	static model = "Xenova/nllb-200-distilled-600M";
@@ -17,7 +15,7 @@ class MyTranslationPipeline {
 			MyTranslationPipeline.instance = pipeline(
 				MyTranslationPipeline.task,
 				MyTranslationPipeline.model,
-				{ progress_callback: progressCallback },
+				{ dtype: "q8", progress_callback: progressCallback },
 			);
 		}
 
@@ -38,7 +36,7 @@ self.addEventListener("message", async (event) => {
 	// Actually perform the translation
 
 	console.log(event.data.tgt_lang, event.data.src_lang, "langcode");
-	const output = await translator(event.data.text, {
+	const output = await translator?.(event.data.text, {
 		tgt_lang: event.data.tgt_lang,
 		src_lang: event.data.src_lang,
 

src/index.ts renamed to src/index.js

@@ -1,6 +1,6 @@
 import renderer from "react-test-renderer";
 import { expect, it } from "vitest";
-import { Example } from "../src";
+import { Example } from "../src/Example";
 
 it("renders correctly", () => {
 	const tree = renderer

src/stories/Example.stories.tsx

@@ -5,16 +5,24 @@ import path from "node:path";
 import { type Options, defineConfig } from "tsup";
 
 const common: Options = {
-	entry: ["src/index.ts"],
+	entry: ["src/index.js", "src/worker.js"],
 	treeshake: false,
 	sourcemap: "inline",
 	minify: true,
 	clean: true,
 	dts: true,
 	splitting: false,
-	format: ["cjs", "esm"],
+	format: ["esm"],
 	external: ["react"],
 	injectStyle: false,
+	target: "es2021",
+	platform: "browser",
+	// dts: {
+	//   entry: 'src/index.js', // or whichever main file you want
+	// },
+	// loader: {
+	//   '.worker.js': 'file', // Handle worker files as static files
+	// },
 };
 
 const getPackageName = async () => {
@@ -66,11 +74,10 @@ const linkSelf = async () => {
 };
 
 export default defineConfig({
-	async onSuccess() {
-		// If you want need to add a use statement to files, you can use the following code:
-		// await _addUseStatement('dist/react', 'client');
+	// async onSuccess() {
+	// 	// If you want need to add a use statement to files, you can use the following code:
+	// 	// await _addUseStatement('dist/react', 'client');
 
-		await linkSelf();
-	},
+	// },
 	...common,
 });