diff --git a/package.json b/package.json index 5545dd5..599bb08 100644 --- a/package.json +++ b/package.json @@ -52,6 +52,7 @@ "license": " AGPL-3.0", "dependencies": { "@dagrejs/dagre": "^1.1.4", + "@huggingface/transformers": "^3.5.2", "@vicons/fluent": "^0.13.0", "@vue-flow/core": "^1.42.5", "better-sqlite3": "^11.10.0", diff --git a/src/main/main.ts b/src/main/main.ts index 013ef35..f930c2d 100644 --- a/src/main/main.ts +++ b/src/main/main.ts @@ -5,6 +5,7 @@ import fs from "fs"; import {spawn} from "child_process"; import './hnswlib-service'; import './database' +import {runEmbedding} from "./model"; // Handle creating/removing shortcuts on Windows when installing/uninstalling. if (started) { @@ -293,3 +294,5 @@ ipcMain.handle('dialog:select-folder', async () => { return result.filePaths[0]; // 返回选中的文件夹路径 } }); + +ipcMain.handle('transformers:runEmbedding', runEmbedding) diff --git a/src/main/model.ts b/src/main/model.ts new file mode 100644 index 0000000..f6a9fff --- /dev/null +++ b/src/main/model.ts @@ -0,0 +1,46 @@ +import {env, FeatureExtractionPipeline, pipeline} from '@huggingface/transformers'; +import {Tensor} from "@huggingface/transformers/types/utils/tensor"; +import IpcMainInvokeEvent = Electron.IpcMainInvokeEvent; + +const EMBEDDING_TASK = "feature-extraction"; + +class EmbeddingsPipeline { + static readonly model: string = 'Qwen3-Embedding-0.6B-ONNX'; + private static instance: FeatureExtractionPipeline | null = null; + + static async getInstance(): Promise { + if (this.instance === null) { + // Dynamically import the Transformers.js library + // Error: Could not dynamically require "../bin/napi-v3/darwin/arm64/onnxruntime_binding.node" + // let { pipeline, env } = await import('@huggingface/transformers'); + + // NOTE: Uncomment this to change the cache directory + // env.cacheDir = './.cache'; + + env.allowLocalModels = true; // 允许使用本地模型 + env.allowRemoteModels = false; // 禁用远程模型下载 + env.localModelPath = '/Users/johnzhang/git-repo/'; // 设置本地模型路径 + + this.instance = await pipeline(EMBEDDING_TASK, this.model, { + // quantized: true, // 明确使用量化模型: model_quantized.onnx + // model_file_name: "model", + dtype: 'q8', + // device: 'gpu' //如果支持 + }); + } + return this.instance; + } +} + +async function runEmbedding(event: IpcMainInvokeEvent, text: string | string[]): Promise { + const extractor = await EmbeddingsPipeline.getInstance(); + const result: Tensor = await extractor(text, {pooling: 'mean', normalize: true}); + console.log(result); + return { + data: result.data, + size: result.size, + location: result.location, + }; +} + +export { runEmbedding }; diff --git a/src/main/preload.ts b/src/main/preload.ts index f7ab35d..c12df93 100644 --- a/src/main/preload.ts +++ b/src/main/preload.ts @@ -45,5 +45,8 @@ contextBridge.exposeInMainWorld('electronAPI', { runQuery: (sql:string, params: any[]) => ipcRenderer.invoke('db:runQuery', sql, params), fetchAll: (sql:string, params: any[]) => ipcRenderer.invoke('db:fetchAll', sql, params), fetchOne: (sql:string, params: any[]) => ipcRenderer.invoke('db:fetchOne', sql, params), + + // Embedding + runEmbedding: (text: string | string[]) => ipcRenderer.invoke('transformers:runEmbedding', text) }); diff --git a/src/renderer/pages/EmbeddingExample.vue b/src/renderer/pages/EmbeddingExample.vue new file mode 100644 index 0000000..042f4a0 --- /dev/null +++ b/src/renderer/pages/EmbeddingExample.vue @@ -0,0 +1,34 @@ + + + + diff --git a/src/renderer/pages/home.vue b/src/renderer/pages/home.vue index 84b136f..103f3d4 100644 --- a/src/renderer/pages/home.vue +++ b/src/renderer/pages/home.vue @@ -20,6 +20,7 @@ play search test scan + Embedding Example diff @@ -102,6 +103,10 @@ const goPlay3= () => { router.push(`/play3`) }; +const tryEmbedding= () => { + router.push(`/embedding-example`) +}; + const handleEditProject= (project) => { // 打开编辑对话框,并预填项目信息 // ... @@ -155,4 +160,4 @@ onMounted(async () => { .body { padding: 20px; } - \ No newline at end of file + diff --git a/src/renderer/router/index.ts b/src/renderer/router/index.ts index c8c52de..9d9750b 100644 --- a/src/renderer/router/index.ts +++ b/src/renderer/router/index.ts @@ -61,6 +61,11 @@ const router = createRouter({ name: 'play3', component: () => import('../pages/task-run-test.vue'), }, + { + path: '/embedding-example', + name: 'embedding-example', + component: () => import('../pages/EmbeddingExample.vue'), + }, { path: '/diff', name: 'diff', diff --git a/src/renderer/types/electron-api.d.ts b/src/renderer/types/electron-api.d.ts index 5b4cb48..b3e32f9 100644 --- a/src/renderer/types/electron-api.d.ts +++ b/src/renderer/types/electron-api.d.ts @@ -23,8 +23,10 @@ interface ElectronAPI { runQuery: (sql: string, params?: any[] | Record) => Promise; fetchAll: (sql: string, params?: any[] | Record) => Promise; fetchOne: (sql: string, params?: any[] | Record) => Promise; + + runEmbedding: (text: string | string[]) => Promise; } declare interface Window { electronAPI: ElectronAPI; -} \ No newline at end of file +} diff --git a/vite.main.config.ts b/vite.main.config.ts index 7fedafe..80db3d1 100644 --- a/vite.main.config.ts +++ b/vite.main.config.ts @@ -3,7 +3,7 @@ import { defineConfig } from 'vite'; export default defineConfig({ build: { rollupOptions: { - external: ['hnswlib-node','better-sqlite3'], // 👈 关键配置:排除原生模块 + external: ['hnswlib-node','better-sqlite3', "@huggingface/transformers"], // 👈 关键配置:排除原生模块 }, }, -}); \ No newline at end of file +});