moving

packages/hive-mind-core/tests/mind/inprocess-reranker.test.ts

@@ -0,0 +1,83 @@
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+const transformers = vi.hoisted(() => ({
+  env: { allowRemoteModels: false, cacheDir: '' },
+  model: vi.fn(),
+  modelFromPretrained: vi.fn(),
+  tokenizer: vi.fn(),
+  tokenizerFromPretrained: vi.fn(),
+}));
+
+vi.mock('@huggingface/transformers', () => ({
+  env: transformers.env,
+  AutoModelForSequenceClassification: {
+    from_pretrained: transformers.modelFromPretrained,
+  },
+  AutoTokenizer: {
+    from_pretrained: transformers.tokenizerFromPretrained,
+  },
+}));
+
+import { createInProcessReranker } from '../../src/mind/inprocess-reranker.js';
+
+const tempRoots: string[] = [];
+
+describe('createInProcessReranker', () => {
+  beforeEach(() => {
+    transformers.env.allowRemoteModels = false;
+    transformers.env.cacheDir = '';
+    transformers.model.mockReset();
+    transformers.modelFromPretrained.mockReset();
+    transformers.tokenizer.mockReset();
+    transformers.tokenizerFromPretrained.mockReset();
+    transformers.modelFromPretrained.mockResolvedValue(transformers.model);
+    transformers.tokenizerFromPretrained.mockResolvedValue(transformers.tokenizer);
+  });
+
+  afterEach(() => {
+    for (const root of tempRoots.splice(0)) {
+      fs.rmSync(root, { force: true, recursive: true });
+    }
+  });
+
+  it('requests tensors and supports single and batch scoring', async () => {
+    transformers.tokenizer.mockResolvedValue({ input_ids: 'tokens' });
+    transformers.model
+      .mockResolvedValueOnce({ logits: { data: new Float32Array([0.75]), dims: [1, 1] } })
+      .mockResolvedValueOnce({ logits: { data: new Float32Array([0.25, 0.5]), dims: [2, 1] } });
+    const cacheDir = fs.mkdtempSync(path.join(os.tmpdir(), 'reranker-test-'));
+    tempRoots.push(cacheDir);
+
+    const reranker = await createInProcessReranker({ cacheDir });
+    await expect(reranker.score('query', 'document')).resolves.toBeCloseTo(0.75);
+    await expect(reranker.scoreBatch('query', ['first', 'second'])).resolves.toEqual([
+      0.25,
+      0.5,
+    ]);
+
+    const canonicalCacheDir = fs.realpathSync.native(cacheDir);
+    expect(transformers.tokenizerFromPretrained).toHaveBeenCalledWith(
+      'Xenova/ms-marco-MiniLM-L-6-v2',
+      { cache_dir: canonicalCacheDir },
+    );
+    expect(transformers.modelFromPretrained).toHaveBeenCalledWith(
+      'Xenova/ms-marco-MiniLM-L-6-v2',
+      { dtype: 'fp32', cache_dir: canonicalCacheDir },
+    );
+    expect(transformers.tokenizer).toHaveBeenNthCalledWith(1, 'query', {
+      text_pair: 'document',
+      padding: true,
+      truncation: true,
+      return_tensor: true,
+    });
+    expect(transformers.tokenizer).toHaveBeenNthCalledWith(2, ['query', 'query'], {
+      text_pair: ['first', 'second'],
+      padding: true,
+      truncation: true,
+      return_tensor: true,
+    });
+  });
+});