Calculating similarity

embeddings/src/emb/data.ts

@@ -4,21 +4,20 @@ import { join } from 'path';
 
 const openai = new OpenAI();
 
-type DataWithEmbeddings = {
+export type DataWithEmbeddings = {
     input: string,
     embedding: number[]
 }
 
-async function generateEmbeddings(input: string| string[]){
+export async function generateEmbeddings(input: string| string[]){
     const response = await openai.embeddings.create({
         input: input,
         model: 'text-embedding-3-small'
     })
-    console.log(response.data[0].embedding)
     return response;
 }
 
-function loadJSONData<T>(fileName: string):T{
+export function loadJSONData<T>(fileName: string):T{
     const path = join(__dirname, fileName);
     const rawData = readFileSync(path);
     return JSON.parse(rawData.toString());

embeddings/src/emb/similar.ts

@@ -0,0 +1,42 @@
+import { DataWithEmbeddings, generateEmbeddings, loadJSONData } from "./data";
+
+function dotProduct(a: number[], b: number[]) {
+    return a.map((value, index) => value * b[index]).reduce((a, b) => a + b, 0);
+}
+
+function cosineSimilarity(a: number[], b: number[]) {
+    const product = dotProduct(a, b);
+    const aMagnitude = Math.sqrt(a.map(value => value * value).reduce((a, b) => a + b, 0));
+    const bMagnitude = Math.sqrt(b.map(value => value * value).reduce((a, b) => a + b, 0));
+    return product / (aMagnitude * bMagnitude);
+}
+
+async function main() {
+    const dataWithEmbeddings = loadJSONData<DataWithEmbeddings[]>('dataWithEmbeddings.json')
+
+    const input = 'animal';
+
+    const inputEmbedding = await generateEmbeddings(input);
+
+    const similarities: {
+        input: string,
+        similarity: number
+    }[] = [];
+
+    for (const entry of dataWithEmbeddings) {
+        const similarity = cosineSimilarity(
+            entry.embedding,
+            inputEmbedding.data[0].embedding
+        )
+        similarities.push({
+            input: entry.input,
+            similarity
+        })
+    }
+
+    console.log(`Similarity of ${input} with:`)
+    const sortedSimilarities = similarities.sort((a, b) => b.similarity - a.similarity);
+    sortedSimilarities.forEach(similarity => {
+        console.log(`${similarity.input}: ${similarity.similarity}`);
+    })
+}