index.js

const path = require('node:path');
const fs = require('node:fs');
const { env, pipeline } = require('@huggingface/transformers');
const { HierarchicalNSW } = require('hnswlib-node');

class BiMap {
  constructor(entries = []) {
    this._forward = new Map(entries);
    this._backward = new Map(entries.map(([k, v]) => [v, k]));
  }
  set(key, value) {
    // 删除已有的正向/反向映射，保证一对一
    if (this._forward.has(key)) {
      const oldVal = this._forward.get(key);
      this._backward.delete(oldVal);
    }
    if (this._backward.has(value)) {
      const oldKey = this._backward.get(value);
      this._forward.delete(oldKey);
    }
    this._forward.set(key, value);
    this._backward.set(value, key);
  }
  get(key) { return this._forward.get(key); }
  getKey(value) { return this._backward.get(value); }
  has(key) { return this._forward.has(key); }
  hasValue(value) { return this._backward.has(value); }
  delete(key) {
    if (!this._forward.has(key)) return false;
    const value = this._forward.get(key);
    this._forward.delete(key);
    this._backward.delete(value);
    return true;
  }
  get length() { return this._forward.size; }
}

const numDimensions = 384; // the length of data point vector that will be indexed.
const maxElements = 1024; // the maximum number of data points.
const modelName = 'Xenova/all-MiniLM-L6-v2';
const config = hexo.config.posts_embedding = Object.assign({
  cache_dir: path.join('.cache', 'hexo-posts-embedding')
}, hexo.config.posts_embedding);

function resolveCacheDir(cacheDir) {
  if (!cacheDir || typeof cacheDir !== 'string') return null;
  if (path.isAbsolute(cacheDir)) return cacheDir;
  return path.join(hexo.base_dir, cacheDir);
}

const cacheDir = resolveCacheDir(config.cache_dir);
if (cacheDir) {
  env.cacheDir = cacheDir;
}

// declaring and intializing index.
const index = new HierarchicalNSW('l2', numDimensions);
index.initIndex(maxElements);

let extractor;
const labelMapping = new BiMap();

function formatBytes(bytes) {
  if (!Number.isFinite(bytes)) return '';
  const units = ['B', 'KB', 'MB', 'GB'];
  let value = bytes;
  let unit = units.shift();
  while (value >= 1024 && units.length) {
    value /= 1024;
    unit = units.shift();
  }
  return `${value.toFixed(value >= 10 || unit === 'B' ? 0 : 1)} ${unit}`;
}

function createProgressLogger(log) {
  const lastLogged = new Map();
  const cachedFiles = new Set();
  let activeDownload;

  return data => {
    if (!data) return;

    const key = data.file
      ? `${data.name || modelName}/${data.file}`
      : activeDownload;

    if (data.status === 'initiate') {
      const cachePath = cacheDir && data.file
        ? path.join(cacheDir, data.name || modelName, data.file)
        : null;
      const isCached = cachePath && fs.existsSync(cachePath);
      if (isCached) cachedFiles.add(key);
      log.info(`${isCached ? 'Loading cached' : 'Loading'} embedding model file: ${key}`);
      return;
    }

    if (data.status === 'download') {
      activeDownload = key;
      return;
    }

    if (data.status === 'done') {
      if (activeDownload === key) activeDownload = null;
      return;
    }

    if (!key) return;
    if (typeof data.progress !== 'number') return;
    if (data.status && data.status !== 'progress') return;
    if (cachedFiles.has(key)) return;

    const progress = Math.floor(data.progress);
    const previous = lastLogged.get(key) || 0;
    if (progress < 100 && progress - previous < 5) return;

    lastLogged.set(key, progress);

    const width = 20;
    const filled = Math.round((Math.min(progress, 100) / 100) * width);
    const bar = `${'#'.repeat(filled)}${'-'.repeat(width - filled)}`;
    const size = data.total
      ? ` (${formatBytes(data.loaded)} / ${formatBytes(data.total)})`
      : '';
    const action = cachedFiles.has(key) ? 'Reading cached' : 'Downloading';

    log.info(`${action} ${key} [${bar}] ${progress}%${size}`);
  };
}

hexo.extend.filter.register('after_init', async function() {
  const log = this.log || hexo.log;
  log.info(`Loading embedding model: ${modelName}`);
  if (cacheDir) {
    log.info(`Embedding model cache: ${cacheDir}`);
  }
  extractor = await pipeline('feature-extraction', modelName, {
    progress_callback: createProgressLogger(log)
  });
  log.info(`Embedding model ready: ${modelName}`);
});

hexo.extend.filter.register('before_post_render', async function(data) {
  const embeddings = await extractor([data._content], { pooling: 'mean', normalize: true });
  data.embedding_vector = embeddings.tolist()[0];
  // Create a new id if data.path doesn't exist in labelMapping
  // Else use the existing id
  let id;
  if (!labelMapping.hasValue(data.path)) {
    id = labelMapping.length;
    labelMapping.set(id, data.path);
  } else {
    id = labelMapping.getKey(data.path);
  }
  index.addPoint(data.embedding_vector, id);
  return data;
});

hexo.extend.helper.register('related_posts', function(post) {
  const result = [];
  if (!post.embedding_vector) {
    post.related_posts = result;
    return result;
  }
  const numNeighbors = 5;
  const query = post.embedding_vector;
  const { neighbors } = index.searchKnn(query, numNeighbors);
  // Skip the first result as it is the query itself
  for (let i = 1; i < neighbors.length; i++) {
    const neighbor = neighbors[i];
    result.push(labelMapping.get(neighbor));
  }
  post.related_posts = result;
  return result;
});