All files / embed-code-node/src ort-backend.ts

100% Statements 22/22
83.33% Branches 5/6
100% Functions 10/10
100% Lines 22/22

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71                  8x             7x                 26x 23x   3x   26x         7x   7x 7x 21x   7x 7x 7x 7x   7x     7x         33x   4x     4x     10x     1x 1x        
/**
 * ORT backend adapter for onnxruntime-node.
 */
import type { IOrtBackend, IOrtSession, IOrtTensor } from '@agentix-e/embed-code-core';
import { int32ToBigInt64 } from '@agentix-e/embed-code-core';
import * as ort from 'onnxruntime-node';
 
export class NodeOrtBackend implements IOrtBackend {
  async createSession(modelPath: string): Promise<IOrtSession> {
    const session = await ort.InferenceSession.create(modelPath, {
      executionProviders: ['cpu'],
      graphOptimizationLevel: 'all',
      enableCpuMemArena: true,
      enableMemPattern: true,
      logSeverityLevel: 3,
    });
    return new NodeOrtSession(session);
  }
 
  createTensor(
    type: 'int64' | 'float32',
    data: number[] | Int32Array | Float32Array | BigInt64Array,
    dims: number[],
  ): IOrtTensor {
    let tensorData: any;
    if (type === 'int64') {
      tensorData = int32ToBigInt64(data as Int32Array | number[]);
    } else {
      tensorData = data;
    }
    return new NodeOrtTensor(new ort.Tensor(type, tensorData, dims));
  }
}
 
class NodeOrtSession implements IOrtSession {
  constructor(private session: ort.InferenceSession) {}
  async run(feeds: Record<string, IOrtTensor>): Promise<Record<string, IOrtTensor>> {
    const feedMap: Record<string, any> = {};
    for (const [k, v] of Object.entries(feeds)) {
      feedMap[k] = (v as NodeOrtTensor).tensor;
    }
    const results = await this.session.run(feedMap);
    const out: Record<string, IOrtTensor> = {};
    for (const [k, v] of Object.entries(results)) {
      out[k] = new NodeOrtTensor(v);
    }
    return out;
  }
  release(): void {
    this.session.release();
  }
}
 
class NodeOrtTensor implements IOrtTensor {
  constructor(public readonly tensor: ort.Tensor) {}
  get type(): string {
    return this.tensor.type;
  }
  get dims(): readonly number[] {
    return this.tensor.dims;
  }
  get data(): Float32Array | Int32Array | BigInt64Array {
    return this.tensor.data as Float32Array | Int32Array | BigInt64Array;
  }
  dispose(): void {
    Eif (this.tensor && typeof (this.tensor as any).dispose === 'function') {
      (this.tensor as any).dispose();
    }
  }
}