Spaces:

NathyB
/

ir-studio

Sleeping

ir-studio / src /utils /indexingEngine.ts

Natnael B. Haile

Cursor

feat: complete IR Studio with Git LFS tracking for all binary assets and NFCorpus index

fbad523 23 days ago

9.56 kB

import type {
  CorpusDocument,
  TermStats,
  DocIndexMeta,
  LengthHistogramBin,
  IndexTelemetry,
  BuiltInvertedIndex,
} from '../types/corpus';
import { computeEmbedding, computeSpladeTokens } from './retrievalScoring';

// Standard Lucene / English Stopwords
export const STANDARD_STOPWORDS = new Set([
  'a', 'about', 'above', 'after', 'again', 'against', 'all', 'am', 'an', 'and', 'any',
  'are', 'aren\'t', 'as', 'at', 'be', 'because', 'been', 'before', 'being', 'below',
  'between', 'both', 'but', 'by', 'can', 'cannot', 'could', 'did', 'do', 'does',
  'doing', 'don\'t', 'down', 'during', 'each', 'few', 'for', 'from', 'further',
  'had', 'has', 'have', 'having', 'he', 'her', 'here', 'hers', 'herself', 'him',
  'himself', 'his', 'how', 'i', 'if', 'in', 'into', 'is', 'it', 'its', 'itself',
  'me', 'more', 'most', 'my', 'myself', 'no', 'nor', 'not', 'of', 'off', 'on',
  'once', 'only', 'or', 'other', 'ought', 'our', 'ours', 'ourselves', 'out', 'over',
  'own', 'same', 'she', 'should', 'so', 'some', 'such', 'than', 'that', 'the',
  'their', 'theirs', 'them', 'themselves', 'then', 'there', 'these', 'they', 'this',
  'those', 'through', 'to', 'too', 'under', 'until', 'up', 'very', 'was', 'we',
  'were', 'what', 'when', 'where', 'which', 'while', 'who', 'whom', 'why', 'with',
  'would', 'you', 'your', 'yours', 'yourself', 'yourselves'
]);

export interface TokenWithPos {
  term: string;
  position: number;
}

// Tokenize text while tracking word position offsets
export function tokenizeWithPositions(
  text: string,
  filterStopwords: boolean = false
): TokenWithPos[] {
  const result: TokenWithPos[] = [];
  const regex = /[a-zA-Z0-9]+(?:'[a-zA-Z0-9]+)?/g;
  let match: RegExpExecArray | null;
  let pos = 0;

  while ((match = regex.exec(text)) !== null) {
    const term = match[0].toLowerCase();
    if (!filterStopwords || !STANDARD_STOPWORDS.has(term)) {
      result.push({ term, position: pos });
    }
    pos++;
  }

  return result;
}

// Compute length histogram bins
export function computeLengthHistogram(docLengths: { id: string; length: number }[]): LengthHistogramBin[] {
  if (docLengths.length === 0) return [];

  const lengths = docLengths.map((d) => d.length);
  const min = Math.min(...lengths);
  const max = Math.max(...lengths);

  if (min === max) {
    return [
      {
        binLabel: `${min} tokens`,
        rangeMin: min,
        rangeMax: max,
        count: docLengths.length,
        percentage: 100,
        docIds: docLengths.map((d) => d.id),
      },
    ];
  }

  const numBins = Math.min(6, Math.max(3, Math.ceil(Math.sqrt(docLengths.length))));
  const binWidth = Math.max(1, Math.ceil((max - min + 1) / numBins));

  const bins: LengthHistogramBin[] = [];
  for (let i = 0; i < numBins; i++) {
    const rangeMin = min + i * binWidth;
    const rangeMax = i === numBins - 1 ? max : rangeMin + binWidth - 1;
    bins.push({
      binLabel: `${rangeMin} - ${rangeMax}`,
      rangeMin,
      rangeMax,
      count: 0,
      percentage: 0,
      docIds: [],
    });
  }

  docLengths.forEach(({ id, length }) => {
    let placed = false;
    for (const bin of bins) {
      if (length >= bin.rangeMin && length <= bin.rangeMax) {
        bin.count++;
        bin.docIds.push(id);
        placed = true;
        break;
      }
    }
    if (!placed && bins.length > 0) {
      bins[bins.length - 1].count++;
      bins[bins.length - 1].docIds.push(id);
    }
  });

  const total = docLengths.length;
  bins.forEach((b) => {
    b.percentage = Number(((b.count / total) * 100).toFixed(1));
  });

  return bins.filter((b) => b.count > 0);
}

export interface IndexBuildOptions {
  filterStopwords?: boolean;
}

/**
 * Builds a complete client-side inverted index for a collection of documents.
 * Emulates Apache Lucene's core in-memory index structures and telemetry.
 */
export function buildInvertedIndex(
  documents: CorpusDocument[],
  options: IndexBuildOptions = {}
): BuiltInvertedIndex {
  const startTime = performance.now();
  const filterStopwords = options.filterStopwords ?? false;

  const N = documents.length;
  const dictionaryMap = new Map<string, TermStats>();
  const docMap = new Map<string, DocIndexMeta>();
  const docIndexList: DocIndexMeta[] = [];

  let totalTokens = 0;
  let rawSizeBytes = 0;

  // 1. Process documents and construct postings
  documents.forEach((doc, internalId) => {
    const fullText = (doc.title ? `${doc.title} ` : '') + doc.text;
    rawSizeBytes += new TextEncoder().encode(fullText).length;

    const tokensWithPos = tokenizeWithPositions(fullText, filterStopwords);
    const docLength = tokensWithPos.length;
    totalTokens += docLength;

    // Per-document term frequency and positions
    const docTfMap = new Map<string, { count: number; positions: number[] }>();
    tokensWithPos.forEach(({ term, position }) => {
      let entry = docTfMap.get(term);
      if (!entry) {
        entry = { count: 0, positions: [] };
        docTfMap.set(term, entry);
      }
      entry.count++;
      entry.positions.push(position);
    });

    // Populate global dictionary & inverted posting lists
    docTfMap.forEach((tfData, term) => {
      let termStat = dictionaryMap.get(term);
      if (!termStat) {
        termStat = {
          term,
          docFreq: 0,
          collectionFreq: 0,
          idf: 0,
          postings: [],
        };
        dictionaryMap.set(term, termStat);
      }

      termStat.docFreq++;
      termStat.collectionFreq += tfData.count;
      termStat.postings.push({
        docId: doc.id,
        docIndex: internalId,
        termFreq: tfData.count,
        positions: tfData.positions,
        docLength,
      });
    });

    // Term vector for document
    const termVector: { term: string; tf: number }[] = [];
    docTfMap.forEach((data, term) => {
      termVector.push({ term, tf: data.count });
    });
    termVector.sort((a, b) => b.tf - a.tf);

    // Compute dense embedding & SPLADE activations
    const denseVector = computeEmbedding(fullText, 768);
    const spladeWeights = computeSpladeTokens(fullText);
    const spladeExpansions = Object.entries(spladeWeights).map(([token, weight]) => ({
      token,
      weight,
      is_expansion: !fullText.toLowerCase().includes(token.toLowerCase()),
    }));

    const docMeta: DocIndexMeta = {
      id: doc.id,
      internalId,
      title: doc.title,
      text: doc.text,
      length: docLength,
      uniqueTerms: docTfMap.size,
      termVector,
      denseVector,
      denseDim: 768,
      denseSource: 'in-browser-simulation',
      denseLatencyMs: 0.1,
      spladeWeights,
      spladeExpansions,
      spladeSource: 'in-browser-simulation',
      spladeLatencyMs: 0.1,
      totalActivatedSplade: Object.keys(spladeWeights).length,
      mlmExpansionsCount: spladeExpansions.filter((e) => e.is_expansion).length,
    };

    docMap.set(doc.id, docMeta);
    docIndexList.push(docMeta);
  });

  // 2. Compute Lucene Robertson-Spärck Jones IDF for every term
  // IDF formula: ln(1 + (N - DF + 0.5) / (DF + 0.5))
  dictionaryMap.forEach((termStat) => {
    const df = termStat.docFreq;
    const rawIdf = Math.log(1.0 + (N - df + 0.5) / (df + 0.5));
    termStat.idf = Number(Math.max(0.1, rawIdf).toFixed(4));
    // Sort postings by Lucene internal Doc ID ascending (standard posting list invariant)
    termStat.postings.sort((a, b) => a.docIndex - b.docIndex);
  });

  const termsList: TermStats[] = Array.from(dictionaryMap.values());
  // Default sort by Document Frequency descending, then Collection Frequency
  termsList.sort((a, b) => {
    if (b.docFreq !== a.docFreq) return b.docFreq - a.docFreq;
    return b.collectionFreq - a.collectionFreq;
  });

  const endTime = performance.now();
  const indexingTimeMs = Number((endTime - startTime).toFixed(2));

  // 3. Compute telemetry metrics
  const avgdl = N > 0 ? Number((totalTokens / N).toFixed(2)) : 0;
  const vocabSize = dictionaryMap.size;

  const docLengths = docIndexList.map((d) => ({ id: d.id, length: d.length }));
  docLengths.sort((a, b) => a.length - b.length);

  const shortestDoc = docLengths[0] || { id: '-', length: 0 };
  const longestDoc = docLengths[docLengths.length - 1] || { id: '-', length: 0 };
  const histogram = computeLengthHistogram(docLengths);

  // Index size estimate:
  // Term dictionary: ~16 bytes overhead per term + string length
  // Postings: 4 bytes docDelta + 2 bytes tf + 2 bytes * positions
  let indexSizeBytes = 0;
  termsList.forEach((t) => {
    indexSizeBytes += 16 + t.term.length; // dictionary entry
    t.postings.forEach((p) => {
      indexSizeBytes += 4 + 2 + (p.positions ? p.positions.length * 2 : 0);
    });
  });

  const compressionRatio =
    indexSizeBytes > 0 ? Number((rawSizeBytes / indexSizeBytes).toFixed(2)) : 1.0;

  const telemetry: IndexTelemetry = {
    docCount: N,
    totalTokens,
    vocabSize,
    avgdl,
    indexingTimeMs,
    rawSizeBytes,
    indexSizeBytes,
    compressionRatio,
    shortestDoc,
    longestDoc,
    histogram,
  };

  return {
    telemetry,
    terms: termsList,
    dictionary: dictionaryMap,
    documents: docIndexList,
    docMap,
  };
}

// Term search filter (prefix or contains)
export function searchIndexTerms(
  index: BuiltInvertedIndex,
  query: string,
  mode: 'prefix' | 'contains' = 'contains'
): TermStats[] {
  if (!query.trim()) return index.terms;
  const q = query.trim().toLowerCase();

  return index.terms.filter((t) => {
    if (mode === 'prefix') {
      return t.term.startsWith(q);
    }
    return t.term.includes(q);
  });
}

	import type {
	CorpusDocument,
	TermStats,
	DocIndexMeta,
	LengthHistogramBin,
	IndexTelemetry,
	BuiltInvertedIndex,
	} from '../types/corpus';
	import { computeEmbedding, computeSpladeTokens } from './retrievalScoring';

	// Standard Lucene / English Stopwords
	export const STANDARD_STOPWORDS = new Set([
	'a', 'about', 'above', 'after', 'again', 'against', 'all', 'am', 'an', 'and', 'any',
	'are', 'aren\'t', 'as', 'at', 'be', 'because', 'been', 'before', 'being', 'below',
	'between', 'both', 'but', 'by', 'can', 'cannot', 'could', 'did', 'do', 'does',
	'doing', 'don\'t', 'down', 'during', 'each', 'few', 'for', 'from', 'further',
	'had', 'has', 'have', 'having', 'he', 'her', 'here', 'hers', 'herself', 'him',
	'himself', 'his', 'how', 'i', 'if', 'in', 'into', 'is', 'it', 'its', 'itself',
	'me', 'more', 'most', 'my', 'myself', 'no', 'nor', 'not', 'of', 'off', 'on',
	'once', 'only', 'or', 'other', 'ought', 'our', 'ours', 'ourselves', 'out', 'over',
	'own', 'same', 'she', 'should', 'so', 'some', 'such', 'than', 'that', 'the',
	'their', 'theirs', 'them', 'themselves', 'then', 'there', 'these', 'they', 'this',
	'those', 'through', 'to', 'too', 'under', 'until', 'up', 'very', 'was', 'we',
	'were', 'what', 'when', 'where', 'which', 'while', 'who', 'whom', 'why', 'with',
	'would', 'you', 'your', 'yours', 'yourself', 'yourselves'
	]);

	export interface TokenWithPos {
	term: string;
	position: number;
	}

	// Tokenize text while tracking word position offsets
	export function tokenizeWithPositions(
	text: string,
	filterStopwords: boolean = false
	): TokenWithPos[] {
	const result: TokenWithPos[] = [];
	const regex = /[a-zA-Z0-9]+(?:'[a-zA-Z0-9]+)?/g;
	let match: RegExpExecArray \| null;
	let pos = 0;

	while ((match = regex.exec(text)) !== null) {
	const term = match[0].toLowerCase();
	if (!filterStopwords \|\| !STANDARD_STOPWORDS.has(term)) {
	result.push({ term, position: pos });
	}
	pos++;
	}

	return result;
	}

	// Compute length histogram bins
	export function computeLengthHistogram(docLengths: { id: string; length: number }[]): LengthHistogramBin[] {
	if (docLengths.length === 0) return [];

	const lengths = docLengths.map((d) => d.length);
	const min = Math.min(...lengths);
	const max = Math.max(...lengths);

	if (min === max) {
	return [
	{
	binLabel: `${min} tokens`,
	rangeMin: min,
	rangeMax: max,
	count: docLengths.length,
	percentage: 100,
	docIds: docLengths.map((d) => d.id),
	},
	];
	}

	const numBins = Math.min(6, Math.max(3, Math.ceil(Math.sqrt(docLengths.length))));
	const binWidth = Math.max(1, Math.ceil((max - min + 1) / numBins));

	const bins: LengthHistogramBin[] = [];
	for (let i = 0; i < numBins; i++) {
	const rangeMin = min + i * binWidth;
	const rangeMax = i === numBins - 1 ? max : rangeMin + binWidth - 1;
	bins.push({
	binLabel: `${rangeMin} - ${rangeMax}`,
	rangeMin,
	rangeMax,
	count: 0,
	percentage: 0,
	docIds: [],
	});
	}

	docLengths.forEach(({ id, length }) => {
	let placed = false;
	for (const bin of bins) {
	if (length >= bin.rangeMin && length <= bin.rangeMax) {
	bin.count++;
	bin.docIds.push(id);
	placed = true;
	break;
	}
	}
	if (!placed && bins.length > 0) {
	bins[bins.length - 1].count++;
	bins[bins.length - 1].docIds.push(id);
	}
	});

	const total = docLengths.length;
	bins.forEach((b) => {
	b.percentage = Number(((b.count / total) * 100).toFixed(1));
	});

	return bins.filter((b) => b.count > 0);
	}

	export interface IndexBuildOptions {
	filterStopwords?: boolean;
	}

	/**
	* Builds a complete client-side inverted index for a collection of documents.
	* Emulates Apache Lucene's core in-memory index structures and telemetry.
	*/
	export function buildInvertedIndex(
	documents: CorpusDocument[],
	options: IndexBuildOptions = {}
	): BuiltInvertedIndex {
	const startTime = performance.now();
	const filterStopwords = options.filterStopwords ?? false;

	const N = documents.length;
	const dictionaryMap = new Map<string, TermStats>();
	const docMap = new Map<string, DocIndexMeta>();
	const docIndexList: DocIndexMeta[] = [];

	let totalTokens = 0;
	let rawSizeBytes = 0;

	// 1. Process documents and construct postings
	documents.forEach((doc, internalId) => {
	const fullText = (doc.title ? `${doc.title} ` : '') + doc.text;
	rawSizeBytes += new TextEncoder().encode(fullText).length;

	const tokensWithPos = tokenizeWithPositions(fullText, filterStopwords);
	const docLength = tokensWithPos.length;
	totalTokens += docLength;

	// Per-document term frequency and positions
	const docTfMap = new Map<string, { count: number; positions: number[] }>();
	tokensWithPos.forEach(({ term, position }) => {
	let entry = docTfMap.get(term);
	if (!entry) {
	entry = { count: 0, positions: [] };
	docTfMap.set(term, entry);
	}
	entry.count++;
	entry.positions.push(position);
	});

	// Populate global dictionary & inverted posting lists
	docTfMap.forEach((tfData, term) => {
	let termStat = dictionaryMap.get(term);
	if (!termStat) {
	termStat = {
	term,
	docFreq: 0,
	collectionFreq: 0,
	idf: 0,
	postings: [],
	};
	dictionaryMap.set(term, termStat);
	}

	termStat.docFreq++;
	termStat.collectionFreq += tfData.count;
	termStat.postings.push({
	docId: doc.id,
	docIndex: internalId,
	termFreq: tfData.count,
	positions: tfData.positions,
	docLength,
	});
	});

	// Term vector for document
	const termVector: { term: string; tf: number }[] = [];
	docTfMap.forEach((data, term) => {
	termVector.push({ term, tf: data.count });
	});
	termVector.sort((a, b) => b.tf - a.tf);

	// Compute dense embedding & SPLADE activations
	const denseVector = computeEmbedding(fullText, 768);
	const spladeWeights = computeSpladeTokens(fullText);
	const spladeExpansions = Object.entries(spladeWeights).map(([token, weight]) => ({
	token,
	weight,
	is_expansion: !fullText.toLowerCase().includes(token.toLowerCase()),
	}));

	const docMeta: DocIndexMeta = {
	id: doc.id,
	internalId,
	title: doc.title,
	text: doc.text,
	length: docLength,
	uniqueTerms: docTfMap.size,
	termVector,
	denseVector,
	denseDim: 768,
	denseSource: 'in-browser-simulation',
	denseLatencyMs: 0.1,
	spladeWeights,
	spladeExpansions,
	spladeSource: 'in-browser-simulation',
	spladeLatencyMs: 0.1,
	totalActivatedSplade: Object.keys(spladeWeights).length,
	mlmExpansionsCount: spladeExpansions.filter((e) => e.is_expansion).length,
	};

	docMap.set(doc.id, docMeta);
	docIndexList.push(docMeta);
	});

	// 2. Compute Lucene Robertson-Spärck Jones IDF for every term
	// IDF formula: ln(1 + (N - DF + 0.5) / (DF + 0.5))
	dictionaryMap.forEach((termStat) => {
	const df = termStat.docFreq;
	const rawIdf = Math.log(1.0 + (N - df + 0.5) / (df + 0.5));
	termStat.idf = Number(Math.max(0.1, rawIdf).toFixed(4));
	// Sort postings by Lucene internal Doc ID ascending (standard posting list invariant)
	termStat.postings.sort((a, b) => a.docIndex - b.docIndex);
	});

	const termsList: TermStats[] = Array.from(dictionaryMap.values());
	// Default sort by Document Frequency descending, then Collection Frequency
	termsList.sort((a, b) => {
	if (b.docFreq !== a.docFreq) return b.docFreq - a.docFreq;
	return b.collectionFreq - a.collectionFreq;
	});

	const endTime = performance.now();
	const indexingTimeMs = Number((endTime - startTime).toFixed(2));

	// 3. Compute telemetry metrics
	const avgdl = N > 0 ? Number((totalTokens / N).toFixed(2)) : 0;
	const vocabSize = dictionaryMap.size;

	const docLengths = docIndexList.map((d) => ({ id: d.id, length: d.length }));
	docLengths.sort((a, b) => a.length - b.length);

	const shortestDoc = docLengths[0] \|\| { id: '-', length: 0 };
	const longestDoc = docLengths[docLengths.length - 1] \|\| { id: '-', length: 0 };
	const histogram = computeLengthHistogram(docLengths);

	// Index size estimate:
	// Term dictionary: ~16 bytes overhead per term + string length
	// Postings: 4 bytes docDelta + 2 bytes tf + 2 bytes * positions
	let indexSizeBytes = 0;
	termsList.forEach((t) => {
	indexSizeBytes += 16 + t.term.length; // dictionary entry
	t.postings.forEach((p) => {
	indexSizeBytes += 4 + 2 + (p.positions ? p.positions.length * 2 : 0);
	});
	});

	const compressionRatio =
	indexSizeBytes > 0 ? Number((rawSizeBytes / indexSizeBytes).toFixed(2)) : 1.0;

	const telemetry: IndexTelemetry = {
	docCount: N,
	totalTokens,
	vocabSize,
	avgdl,
	indexingTimeMs,
	rawSizeBytes,
	indexSizeBytes,
	compressionRatio,
	shortestDoc,
	longestDoc,
	histogram,
	};

	return {
	telemetry,
	terms: termsList,
	dictionary: dictionaryMap,
	documents: docIndexList,
	docMap,
	};
	}

	// Term search filter (prefix or contains)
	export function searchIndexTerms(
	index: BuiltInvertedIndex,
	query: string,
	mode: 'prefix' \| 'contains' = 'contains'
	): TermStats[] {
	if (!query.trim()) return index.terms;
	const q = query.trim().toLowerCase();

	return index.terms.filter((t) => {
	if (mode === 'prefix') {
	return t.term.startsWith(q);
	}
	return t.term.includes(q);
	});
	}