1
>
/*---------------------------------------------------------------------------------------------
tfIdf.ts
2
>
* Copyright (c) Microsoft Corporation. All rights reserved.
3
>
* Licensed under the MIT License. See License.txt in the project root for license information.
4
>
*--------------------------------------------------------------------------------------------*/
5
>
6
>
import { CancellationToken } from './cancellation.js';
7
>
8
>
type SparseEmbedding = Record</* word */ string, /* weight */number>;
9
>
type TermFrequencies = Map</* word */ string, /*occurrences*/ number>;
10
>
type DocumentOccurrences = Map</* word */ string, /*documentOccurrences*/ number>;
11
>
12
>
function countMapFrom<K>(values: Iterable<K>): Map<K, number> {
13
>
const map = new Map<K, number>();
14
>
for (const value of values) {
15
>
map.set(value, (map.get(value) ?? 0) + 1);
16
>
}
17
>
return map;
18
>
}
19
>
20
>
interface DocumentChunkEntry {
21
>
readonly text: string;
22
>
readonly tf: TermFrequencies;
23
>
}
24
>
25
>
export interface TfIdfDocument {
26
>
readonly key: string;
27
>
readonly textChunks: readonly string[];
28
>
}
29
>
30
>
export interface TfIdfScore {
31
>
readonly key: string;
32
>
/**
33
>
* An unbounded number.
34
>
*/
35
>
readonly score: number;
36
>
}
37
>
38
>
export interface NormalizedTfIdfScore {
39
>
readonly key: string;
40
>
/**
41
>
* A number between 0 and 1.
42
>
*/
43
>
readonly score: number;
44
>
}
45
>
46
>
/**
47
>
* Implementation of tf-idf (term frequency-inverse document frequency) for a set of
48
>
* documents where each document contains one or more chunks of text.
49
>
* Each document is identified by a key, and the score for each document is computed
50
>
* by taking the max score over all the chunks in the document.
51
>
*/
52
>
export class TfIdfCalculator {
53
>
calculateScores(query: string, token: CancellationToken): TfIdfScore[] {
54
>
const embedding = this.computeEmbedding(query);
55
>
const idfCache = new Map<string, number>();
56
>
const scores: TfIdfScore[] = [];
57
>
// For each document, generate one score
58
>
for (const [key, doc] of this.documents) {
59
if (token.isCancellationRequested) {
60
return [];