modified faiss

model.py

@@ -75,16 +75,12 @@ def faiss_add_index_cos(df, column):
           
     # Create an index
     index = faiss.IndexFlatIP(embeddings.shape[1])
-    print("<<<<faiss_ after normalize") 
-    faiss.normalize_L2(embeddings)
-    print("<<<<faiss_ after normalize") 
+    faiss.normalize_L2(embeddings) 
 
     index.train(embeddings)
-    print("<<<<faiss_ after index.train")
 
     # Add the embeddings to the index
     index.add(embeddings)
-    print("<<<<faiss_add")
 
     # Return the index
     return index
@@ -100,7 +96,7 @@ def faiss_get_top_N_images(query,
                                           model, tokenizer, 
                                           device)
     # Relevant columns
-    relevant_cols = ["comment", "image_name", "similarity"]
+    relevant_cols = ["comment", "image_name"]
 
     #faiss search with cos similarity
     index = faiss_add_index_cos(data, column="text_embeddings")
@@ -113,5 +109,7 @@ def faiss_get_top_N_images(query,
     non_repeated_images = ~data_sorted["image_name"].duplicated()
     most_similar_articles = data_sorted[non_repeated_images].head(top_K)
 
-    result_df = most_similar_articles[relevant_cols].reset_index(), D.reshape(-1,1)[:top_K]
+    result_df = most_similar_articles[relevant_cols].reset_index()
+    D = D.reshape(-1,1)[:top_K]
+    result_df = pd.concat([result_df, pd.DataFrame(D, columns=['similarity'])], axis=1)
     return [get_item_data(result_df, i, 'similarity') for i in range(len(result_df))]