add device option (cpu / webgpu)

Author: jparkerweb

chunkit.js

@@ -40,6 +40,7 @@ export async function chunkit(
         combineChunksSimilarityThreshold = DEFAULT_CONFIG.COMBINE_CHUNKS_SIMILARITY_THRESHOLD,
         onnxEmbeddingModel = DEFAULT_CONFIG.ONNX_EMBEDDING_MODEL,
         dtype = DEFAULT_CONFIG.DTYPE,
+        device = DEFAULT_CONFIG.DEVICE,
         localModelPath = DEFAULT_CONFIG.LOCAL_MODEL_PATH,
         modelCacheDir = DEFAULT_CONFIG.MODEL_CACHE_DIR,
         returnEmbedding = DEFAULT_CONFIG.RETURN_EMBEDDING,
@@ -59,6 +60,7 @@ export async function chunkit(
     const { modelName, dtype: usedDtype } = await initializeEmbeddingUtils(
         onnxEmbeddingModel, 
         dtype,
+        device,
         localModelPath,
         modelCacheDir
     );
@@ -187,6 +189,7 @@ export async function cramit(
         maxTokenSize = DEFAULT_CONFIG.MAX_TOKEN_SIZE,
         onnxEmbeddingModel = DEFAULT_CONFIG.ONNX_EMBEDDING_MODEL,
         dtype = DEFAULT_CONFIG.DTYPE,
+        device = DEFAULT_CONFIG.DEVICE,
         localModelPath = DEFAULT_CONFIG.LOCAL_MODEL_PATH,
         modelCacheDir = DEFAULT_CONFIG.MODEL_CACHE_DIR,
         returnEmbedding = DEFAULT_CONFIG.RETURN_EMBEDDING,
@@ -206,6 +209,7 @@ export async function cramit(
     await initializeEmbeddingUtils(
         onnxEmbeddingModel, 
         dtype,
+        device,
         localModelPath,
         modelCacheDir
     );
@@ -293,6 +297,7 @@ export async function sentenceit(
         logging = DEFAULT_CONFIG.LOGGING,
         onnxEmbeddingModel = DEFAULT_CONFIG.ONNX_EMBEDDING_MODEL,
         dtype = DEFAULT_CONFIG.DTYPE,
+        device = DEFAULT_CONFIG.DEVICE,
         localModelPath = DEFAULT_CONFIG.LOCAL_MODEL_PATH,
         modelCacheDir = DEFAULT_CONFIG.MODEL_CACHE_DIR,
         returnEmbedding = DEFAULT_CONFIG.RETURN_EMBEDDING,
@@ -313,6 +318,7 @@ export async function sentenceit(
         await initializeEmbeddingUtils(
             onnxEmbeddingModel, 
             dtype,
+            device,
             localModelPath,
             modelCacheDir
         );

config.js

@@ -9,6 +9,7 @@ export const DEFAULT_CONFIG = {
     COMBINE_CHUNKS_SIMILARITY_THRESHOLD: 0.5,
     ONNX_EMBEDDING_MODEL: "Xenova/all-MiniLM-L6-v2",
     DTYPE: 'q8',
+    DEVICE: 'cpu', // webgpu
     LOCAL_MODEL_PATH: "./models",
     MODEL_CACHE_DIR: "./models",
     RETURN_EMBEDDING: false,

embeddingUtils.js

@@ -18,6 +18,7 @@ const embeddingCache = new LRUCache({
 export async function initializeEmbeddingUtils(
     onnxEmbeddingModel, 
     dtype = 'fp32',
+    device = 'cpu',
     localModelPath = null,
     modelCacheDir = null
 ) {
@@ -27,15 +28,22 @@ export async function initializeEmbeddingUtils(
     if (modelCacheDir) env.cacheDir = modelCacheDir;
 
     tokenizer = await AutoTokenizer.from_pretrained(onnxEmbeddingModel);
-    generateEmbedding = await pipeline('feature-extraction', onnxEmbeddingModel, {
+    const pipelineOptions = {
         dtype: dtype,
-    });
+    };
+
+    if (device !== 'webgpu') {
+        pipelineOptions.device = device;
+    }
+
+    generateEmbedding = await pipeline('feature-extraction', onnxEmbeddingModel, pipelineOptions);
 
     embeddingCache.clear();
 
     return {
         modelName: onnxEmbeddingModel,
-        dtype: dtype
+        dtype: dtype,
+        device: device,
     };
 }
 

example/example-chunkit.js

@@ -31,6 +31,9 @@ for (const textFile of textFiles) {
     });
 }
 
+// Get device from command line arguments, default to 'cpu'
+const device = process.argv[2] || 'cpu';
+
 // start timing
 const startTime = performance.now();
 
@@ -47,6 +50,7 @@ let myTestChunks = await chunkit(
         combineChunksSimilarityThreshold: 0.700,
         onnxEmbeddingModel: "Xenova/all-MiniLM-L6-v2",
         dtype: "q8",
+        device: device, // Pass the device to chunkit
         localModelPath: "../models",
         modelCacheDir: "../models",
         returnTokenLength: true,
@@ -61,8 +65,8 @@ const endTime = performance.now();
 let trackedTimeSeconds = (endTime - startTime) / 1000;
 trackedTimeSeconds =  parseFloat(trackedTimeSeconds.toFixed(2));
 
-console.log("\n\n");
 console.log("myTestChunks:");
 console.log(myTestChunks);
+console.log(`device: ${device}`);
 console.log("length: " + myTestChunks.length);
 console.log("trackedTimeSeconds: " + trackedTimeSeconds);

example/example-sentenceit.js

@@ -34,6 +34,9 @@ for (const textFile of textFiles) {
     });
 }
 
+// Get device from command line arguments, default to 'cpu'
+const device = process.argv[2] || 'cpu';
+
 // start timing
 const startTime = performance.now();
 
@@ -43,6 +46,7 @@ let myTestSentences = await sentenceit(
         logging: false,
         onnxEmbeddingModel: "Xenova/all-MiniLM-L6-v2",
         dtype: 'fp32',
+        device: device,
         localModelPath: "../models",
         modelCacheDir: "../models",
         returnEmbedding: true,
@@ -59,5 +63,6 @@ trackedTimeSeconds =  parseFloat(trackedTimeSeconds.toFixed(2));
 console.log("\n\n\n");
 console.log("myTestSentences:");
 console.log(myTestSentences);
+console.log(`device: ${device}`);
 console.log("length: " + myTestSentences.length);
-console.log("trackedTimeSeconds: " + trackedTimeSeconds);
+console.log("trackedTimeSeconds: " + trackedTimeSeconds);

webui/README.md

@@ -1,109 +1,110 @@
-# 🍱 Semantic Chunking Web UI
-
-A web-based interface for experimenting with and tuning Semantic Chunking settings. This tool provides a visual way to test and configure the `semantic-chunking` library's settings to get optimal results for your specific use case. Once you've found the best settings, you can generate code to implement them in your project.
-
-## Features
-
-- Real-time text chunking with live preview
-- Interactive controls for all chunking parameters
-- Visual feedback for similarity thresholds
-- Model selection and configuration
-- Results download in JSON format
-- Code generation for your settings
-- Example texts for testing
-- Dark mode interface
-- Syntax highlighting of JSON results and code samples
-- Line wrapping toggle for JSON results
-
-![semantic-chunking_web-ui](../img/semantic-chunking_web-ui.gif)
-
-## Getting Started
-
-### Prerequisites
-- Node.js (v18 or higher recommended)
-- npm (comes with Node.js)
-
-### Installation
-
-1. Clone the repository: 
-```bash
-git clone https://github.com/jparkerweb/semantic-chunking.git
-```
-
-2. Navigate to the webui directory:
-```bash
-cd semantic-chunking/webui
-```
-
-3. Install dependencies:
-```bash
-npm install
-```
-
-4. Start the server:
-```bash
-npm start
-```
-
-5. Open your browser and visit:
-```bash
-http://localhost:3000
-```
-
-## Usage
-
-### Basic Controls
-
-- **Document Name**: Name for your input text
-- **Text to Chunk**: Your input text to be processed
-- **Max Token Size**: Maximum size for each chunk (50-2500 tokens)
-- **Similarity Threshold**: Base threshold for semantic similarity (0.1-1.0)
-- **Similarity Sentences Lookahead**: Number of sentences to look ahead when calculating similarity (1-10)
-
-### Advanced Settings
-
-- **Dynamic Threshold Bounds**: Lower and upper bounds for dynamic similarity threshold adjustment
-- **Combine Chunks**: Enable/disable chunk combination phase
-- **Combine Chunks Similarity Threshold**: Threshold for combining similar chunks
-
-### Model Settings
-
-- **Embedding Model**: Choose from various supported embedding models
-- **Quantized Model**: Toggle model quantization for reduced memory usage
-
-### Output Settings
-
-- **Return Token Length**: Include token count in results
-- **Return Embedding**: Include embeddings in results
-- **Chunk Prefix**: Add prefix to chunks (useful for RAG applications)
-- **Exclude Chunk Prefix in Results**: Remove prefix from final results
-
-### Example Texts
-
-Use the provided example texts to test different scenarios:
-- `similar.txt`: Text with high semantic similarity between sentences
-- `different.txt`: Text with low semantic similarity between sentences
-
-### Results
-
-- View chunked results in real-time
-- See chunk count, average token length, and processing time
-- Download results as JSON
-- Get generated code with your current settings
-
-## Development
-
-The web UI is built with:
-- `semantic-chunking` library for text processing
-- Express.js for the backend
-- Vanilla JavaScript (ES6+) for the frontend
-- CSS3 for styling
-
-## License
-
-This project is licensed under the MIT License - see the LICENSE file for details.
-
-## Appreciation
-
-If you enjoy this package please consider sending me a tip to support my work 😀
-# [🍵 tip me here](https://ko-fi.com/jparkerweb)
+# 🍱 Semantic Chunking Web UI
+
+A web-based interface for experimenting with and tuning Semantic Chunking settings. This tool provides a visual way to test and configure the `semantic-chunking` library's settings to get optimal results for your specific use case. Once you've found the best settings, you can generate code to implement them in your project.
+
+## Features
+
+- Real-time text chunking with live preview
+- Interactive controls for all chunking parameters
+- Visual feedback for similarity thresholds
+- Model selection and configuration
+- Results download in JSON format
+- Code generation for your settings
+- Example texts for testing
+- Dark mode interface
+- Syntax highlighting of JSON results and code samples
+- Line wrapping toggle for JSON results
+
+![semantic-chunking_web-ui](../img/semantic-chunking_web-ui.gif)
+
+## Getting Started
+
+### Prerequisites
+- Node.js (v18 or higher recommended)
+- npm (comes with Node.js)
+
+### Installation
+
+1. Clone the repository: 
+```bash
+git clone https://github.com/jparkerweb/semantic-chunking.git
+```
+
+2. Navigate to the webui directory:
+```bash
+cd semantic-chunking/webui
+```
+
+3. Install dependencies:
+```bash
+npm install
+```
+
+4. Start the server:
+```bash
+npm start
+```
+
+5. Open your browser and visit:
+```bash
+http://localhost:3000
+```
+
+## Usage
+
+### Basic Controls
+
+- **Document Name**: Name for your input text
+- **Text to Chunk**: Your input text to be processed
+- **Max Token Size**: Maximum size for each chunk (50-2500 tokens)
+- **Similarity Threshold**: Base threshold for semantic similarity (0.1-1.0)
+- **Similarity Sentences Lookahead**: Number of sentences to look ahead when calculating similarity (1-10)
+
+### Advanced Settings
+
+- **Dynamic Threshold Bounds**: Lower and upper bounds for dynamic similarity threshold adjustment
+- **Combine Chunks**: Enable/disable chunk combination phase
+- **Combine Chunks Similarity Threshold**: Threshold for combining similar chunks
+
+### Model Settings
+
+- **Embedding Model**: Choose from various supported embedding models
+- **DType**: Select the data type for the model, affecting precision and performance (e.g., `fp32`, `fp16`, `q8`).
+- **Device**: Choose the processing device (`cpu` or `webgpu`).
+
+### Output Settings
+
+- **Return Token Length**: Include token count in results
+- **Return Embedding**: Include embeddings in results
+- **Chunk Prefix**: Add prefix to chunks (useful for RAG applications)
+- **Exclude Chunk Prefix in Results**: Remove prefix from final results
+
+### Example Texts
+
+Use the provided example texts to test different scenarios:
+- `similar.txt`: Text with high semantic similarity between sentences
+- `different.txt`: Text with low semantic similarity between sentences
+
+### Results
+
+- View chunked results in real-time
+- See chunk count, average token length, and processing time
+- Download results as JSON
+- Get generated code with your current settings
+
+## Development
+
+The web UI is built with:
+- `semantic-chunking` library for text processing
+- Express.js for the backend
+- Vanilla JavaScript (ES6+) for the frontend
+- CSS3 for styling
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.
+
+## Appreciation
+
+If you enjoy this package please consider sending me a tip to support my work 😀
+# [🍵 tip me here](https://ko-fi.com/jparkerweb)