cosine similarity

query-cosine.ts

/**
 * Run cosine similarity search on all documents in the database against a list of queries, and return the top {numResults} results
 * @param documentEmbeddings number[][] - embeddings of source document
 * @param queryEmbeddings number[][] - embeddings of queries (queries are a list)
 * @param numResults number - number of results to be returned
 * @returns SearchResult[][] - list of results
 */
export function search(documentEmbeddings: number[][], queryEmbeddings: number[][], numResults: number = 3): SearchResults[][] {
    const results: SearchResults[][] = [];
    queryEmbeddings.map((queryEmbedding) => {
        const similarityRankings = documentEmbeddings.map((vector, i) => {
            const similarity = cosineSimilarity(vector, queryEmbedding);
            return { 
                index: i,
                similarity,
                vector
            }
        })
        .sort((a, b) => b.similarity - a.similarity);
        
        const slicedRankings = similarityRankings.slice(0, numResults);
        results.push(slicedRankings)
    });
    return results;
}