Upload book_metadata_retriever.py

book_metadata_retriever.py

@@ -1,16 +1,10 @@
-
-#pip install rank-bm25
-
 import numpy as np
 import pandas as pd
 from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.metrics.pairwise import cosine_similarity
-from sklearn.model_selection import train_test_split
 from rank_bm25 import BM25Okapi
 
-
 # Read CSV file
-data = pd.read_csv('books.csv', encoding='latin1')
+data = pd.read_csv(r'C:\book_metadata_retriever\books.csv', encoding='latin1')
 
 class TFIDFDoc2Vec:
     def __init__(self):
@@ -21,59 +15,60 @@ class TFIDFDoc2Vec:
         tfidf_matrix = self.tfidf_vectorizer.fit_transform(documents)
         self.doc_vectors = tfidf_matrix.toarray()
 
-    def find_similar_documents(self, query, top_n=10):
+    def find_similar_documents(self, query, threshold=0.5):
         query_vector = self.tfidf_vectorizer.transform([query]).toarray()
-        similarities = cosine_similarity(query_vector, self.doc_vectors)
-        similar_indices = similarities.argsort(axis=1)[:, ::-1][:, :top_n]
+        similarities = np.dot(query_vector, self.doc_vectors.T)
+        similar_indices = np.where(similarities >= threshold)[1]
+        return similar_indices, similarities
 
-        similar_documents = []
-        for indices in similar_indices:
-            similar_documents.append(indices)
-        return similar_documents
+def answer(query, threshold=0.5, top_n=10):
+    # Find similar documents using TF-IDF
+    similar_documents_indices, similarities = tfidf_doc2vec_model.find_similar_documents(query, threshold=threshold)
 
-    def rank_bm25(self, query, bm25_model, documents, top_n=10):
-        scores = bm25_model.get_scores(query)
-        top_indices = np.argsort(scores)[::-1][:top_n]
-        return top_indices
-data
+    # Check if no similar documents are found
+    if len(similar_documents_indices) == 0:
+        return "No books found for the query."
 
-# Select the column containing book titles
-documents = data['Book Title'].astype(str)
+    # Rank similar documents using BM25
+    scores = bm25_model.get_scores(query.split())  # Split the query into tokens
+    bm25_ranked_indices = np.argsort(scores)[::-1]
+
+    # Initialize a set to keep track of unique document indices
+    unique_indices = set()
+
+    # Combine results from TF-IDF and BM25, keeping unique indices
+    combined_indices = []
+    for index in similar_documents_indices:
+        if index not in unique_indices:
+            combined_indices.append(index)
+            unique_indices.add(index)
+    for index in bm25_ranked_indices:
+        if index not in unique_indices:
+            combined_indices.append(index)
+            unique_indices.add(index)
+
+    # Retrieve document details
+    ranked_documents = []
+    for index in combined_indices[:top_n]:  # Adjust to the desired number of results
+        document_details = {
+            "Book": data['Book Title'].iloc[index],
+            "Author": data['Author'].iloc[index],
+            "Edition": data['Edition'].iloc[index],
+            "File Name": data['File_name'].iloc[index]
+        }
+        ranked_documents.append(document_details)
+
+    return ranked_documents
 
-# Initialize TF-IDF vectors and model
+# Initialize TF-IDF model
 tfidf_doc2vec_model = TFIDFDoc2Vec()
+documents = data['Book Title'].astype(str)
 tfidf_doc2vec_model.initialize_vectors(documents)
 
 # Initialize BM25 model
-bm25_model = BM25Okapi(documents.str.split())
-
-
-def answer(query):
-    # Find similar documents
-    similar_documents_indices = tfidf_doc2vec_model.find_similar_documents(query)
-
-    # Rank similar documents using BM25
-    similar_documents_indices_bm25 = tfidf_doc2vec_model.rank_bm25(query, bm25_model, documents)
-
-    # Initialize a list to store ranked documents
-    ranked_documents = []
-
-    # Add details of each document to the list
-    for indices in similar_documents_indices:
-        for index in indices:
-            document_details = {
-                "Book": data['Book Title'].iloc[index],
-                "Author": data['Author'].iloc[index],
-                "Copyright Year": data['Copyright Year'].iloc[index],
-                "Edition": data['Edition'].iloc[index],
-                "File Name": data['File_name'].iloc[index]
-                
-            }
-            ranked_documents.append(document_details)
-
-    return ranked_documents
+bm25_model = BM25Okapi([doc.split() for doc in documents])
 
-# Receive query from the user
-#query = input("Enter your query: ")
-#result = answer(query)
-#print(result)
+# Example usage
+query = "mathematics"
+result = answer(query)
+print(result)