Kaynağa Gözat

adds computeTextEmbeddings API

Abhinav 1 yıl önce
ebeveyn
işleme
43c806be76
4 değiştirilmiş dosya ile 63 ekleme ve 2 silme
  1. 10 0
      src/api/clip.ts
  2. 2 1
      src/preload.ts
  3. 44 0
      src/services/clipService.ts
  4. 7 1
      src/utils/ipcComms.ts

+ 10 - 0
src/api/clip.ts

@@ -20,3 +20,13 @@ export async function computeImageEmbeddings(
         }
     }
 }
+
+export async function computeTextEmbeddings(
+    text: string
+): Promise<Float32Array> {
+    const embeddings = await ipcRenderer.invoke(
+        'compute-text-embeddings',
+        text
+    );
+    return embeddings;
+}

+ 2 - 1
src/preload.ts

@@ -59,7 +59,7 @@ import {
     logRendererProcessMemoryUsage,
 } from './utils/processStats';
 import { runFFmpegCmd } from './api/ffmpeg';
-import { computeImageEmbeddings } from './api/clip';
+import { computeImageEmbeddings, computeTextEmbeddings } from './api/clip';
 
 fixHotReloadNext12();
 setupLogging();
@@ -116,4 +116,5 @@ windowObject['ElectronAPIs'] = {
     deleteFile,
     updateOptOutOfCrashReports,
     computeImageEmbeddings,
+    computeTextEmbeddings,
 };

+ 44 - 0
src/services/clipService.ts

@@ -16,6 +16,14 @@ const IMAGE_EMBEDDING_EXTRACT_CMD: string[] = [
     INPUT_PATH_PLACEHOLDER,
 ];
 
+const TEXT_EMBEDDING_EXTRACT_CMD: string[] = [
+    GGMLCLIP_PATH_PLACEHOLDER,
+    '-m',
+    CLIP_MODEL_PATH_PLACEHOLDER,
+    '--text',
+    INPUT_PATH_PLACEHOLDER,
+];
+
 function getClipModelPath() {
     return './models/openai_clip-vit-base-patch32.ggmlv0.f16.bin';
 }
@@ -58,3 +66,39 @@ export async function computeImageEmbeddings(
         logErrorSentry(err, 'Error in computeImageEmbeddings');
     }
 }
+
+export async function computeTextEmbeddings(
+    text: string
+): Promise<Float32Array> {
+    try {
+        const clipModelPath = getClipModelPath();
+        const ggmlclipPath = getGGMLClipPath();
+        const cmd = TEXT_EMBEDDING_EXTRACT_CMD.map((cmdPart) => {
+            if (cmdPart === GGMLCLIP_PATH_PLACEHOLDER) {
+                return ggmlclipPath;
+            } else if (cmdPart === CLIP_MODEL_PATH_PLACEHOLDER) {
+                return clipModelPath;
+            } else if (cmdPart === INPUT_PATH_PLACEHOLDER) {
+                return text;
+            } else {
+                return cmdPart;
+            }
+        });
+
+        const escapedCmd = shellescape(cmd);
+        log.info('running clip command', escapedCmd);
+        const startTime = Date.now();
+        const { stdout } = await execAsync(escapedCmd);
+        log.info('clip command execution time ', Date.now() - startTime);
+        // parse stdout and return embeddings
+        // get the last line of stdout
+        const lines = stdout.split('\n');
+        const lastLine = lines[lines.length - 1];
+        const embeddings = JSON.parse(lastLine);
+        const embeddingsArray = new Float32Array(embeddings);
+        console.log('embeddingsArray', embeddingsArray);
+        return embeddingsArray;
+    } catch (err) {
+        logErrorSentry(err, 'Error in computeImageEmbeddings');
+    }
+}

+ 7 - 1
src/utils/ipcComms.ts

@@ -27,7 +27,10 @@ import { deleteTempFile, runFFmpegCmd } from '../services/ffmpeg';
 import { generateTempFilePath } from './temp';
 import { setOptOutOfCrashReports } from '../services/userPreference';
 import { updateOptOutOfCrashReports } from '../main';
-import { computeImageEmbeddings } from '../services/clipService';
+import {
+    computeImageEmbeddings,
+    computeTextEmbeddings,
+} from '../services/clipService';
 
 export default function setupIpcComs(
     tray: Tray,
@@ -170,4 +173,7 @@ export default function setupIpcComs(
     ipcMain.handle('compute-image-embeddings', (_, inputFilePath) => {
         return computeImageEmbeddings(inputFilePath);
     });
+    ipcMain.handle('compute-text-embeddings', (_, text) => {
+        return computeTextEmbeddings(text);
+    });
 }