Skip to content

Instantly share code, notes, and snippets.

@bramses
Created Jan 4, 2022
Embed
What would you like to do?
cosine similarity
/**
* 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;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment