Export LLMPipeline directly

openvinotoolkit · Jan 29, 2025 · 21fdf7b · 21fdf7b
1 parent a6f8b55
commit 21fdf7b
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 16 deletions.
diff --git a/samples/js/text_generation/chat_sample.js b/samples/js/text_generation/chat_sample.js
@@ -1,5 +1,5 @@
 import readline from 'readline';
-import { Pipeline } from 'genai-node';
+import { LLMPipeline } from 'genai-node';
 
 main();
 
@@ -24,7 +24,7 @@ async function main() {
     output: process.stdout,
   });
 
-  const pipe = await Pipeline.LLMPipeline(MODEL_PATH, device);
+  const pipe = await LLMPipeline(MODEL_PATH, device);
   const config = { 'max_new_tokens': 100 };
 
   await pipe.startChat();

diff --git a/src/js/README.md b/src/js/README.md
@@ -13,9 +13,9 @@ npm install genai-node
 
 Use the **genai-node** package:
 ```js
-import { Pipeline } from 'genai-node';
+import { LLMPipeline } from 'genai-node';
 
-const pipe = await Pipeline.LLMPipeline(MODEL_PATH, device);
+const pipe = await LLMPipeline(MODEL_PATH, device);
 
 const input = 'What is the meaning of life?';
 const config = { 'max_new_tokens': 100 };

diff --git a/src/js/lib/module.js b/src/js/lib/module.js
@@ -125,7 +125,7 @@ class LLMPipeline {
   }
 }
 
-class Pipeline {
+class PipelineFactory {
   static async LLMPipeline(modelPath, device = 'CPU') {
     const pipeline = new LLMPipeline(modelPath, device);
     await pipeline.init();
@@ -135,7 +135,5 @@ class Pipeline {
 }
 
 
-export {
-  addon,
-  Pipeline,
-};
+export { addon };
+export const LLMPipeline = PipelineFactory.LLMPipeline;
diff --git a/src/js/tests/module.test.js b/src/js/tests/module.test.js
@@ -1,4 +1,4 @@
-import { Pipeline } from '../lib/module.js';
+import { LLMPipeline } from '../lib/module.js';
 
 import assert from 'node:assert/strict';
 import { describe, it, before, after } from 'node:test';
@@ -11,7 +11,7 @@ describe('module', async () => {
   let pipeline = null;
 
   await before(async () => {
-    pipeline = await Pipeline.LLMPipeline(MODEL_PATH, 'CPU');
+    pipeline = await LLMPipeline(MODEL_PATH, 'CPU');
 
     await pipeline.startChat();
   });
@@ -33,7 +33,7 @@ describe('module', async () => {
 
 describe('corner cases', async () => {
   it('should throw an error if pipeline is already initialized', async () => {
-    const pipeline = await Pipeline.LLMPipeline(MODEL_PATH, 'CPU');
+    const pipeline = await LLMPipeline(MODEL_PATH, 'CPU');
 
     await assert.rejects(
       async () => await pipeline.init(),
@@ -45,7 +45,7 @@ describe('corner cases', async () => {
   });
 
   it('should throw an error if chat is already started', async () => {
-    const pipeline = await Pipeline.LLMPipeline(MODEL_PATH, 'CPU');
+    const pipeline = await LLMPipeline(MODEL_PATH, 'CPU');
 
     await pipeline.startChat();
 
@@ -59,7 +59,7 @@ describe('corner cases', async () => {
   });
 
   it('should throw an error if chat is not started', async () => {
-    const pipeline = await Pipeline.LLMPipeline(MODEL_PATH, 'CPU');
+    const pipeline = await LLMPipeline(MODEL_PATH, 'CPU');
 
     await assert.rejects(
       () => pipeline.finishChat(),
@@ -75,7 +75,7 @@ describe('generation parameters validation', () => {
   let pipeline = null;
 
   before(async () => {
-    pipeline = await Pipeline.LLMPipeline(MODEL_PATH, 'CPU');
+    pipeline = await LLMPipeline(MODEL_PATH, 'CPU');
 
     await pipeline.startChat();
   });
@@ -95,7 +95,7 @@ describe('generation parameters validation', () => {
   });
 
   it('should throw an error if generationCallback is not a function', async () => {
-    const pipeline = await Pipeline.LLMPipeline(MODEL_PATH, 'CPU');
+    const pipeline = await LLMPipeline(MODEL_PATH, 'CPU');
 
     await pipeline.startChat();