book_metadata_retriever / book_metadata_retriever.py
achdaisy's picture
Update book_metadata_retriever.py
00f3159 verified
raw history blame
No virus
2.47 kB
#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')
class TFIDFDoc2Vec:
def __init__(self):
self.tfidf_vectorizer = TfidfVectorizer()
self.doc_vectors = None
def initialize_vectors(self, documents):
tfidf_matrix = self.tfidf_vectorizer.fit_transform(documents)
self.doc_vectors = tfidf_matrix.toarray()
def find_similar_documents(self, query, top_n=10):
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]
similar_documents = []
for indices in similar_indices:
similar_documents.append(indices)
return similar_documents
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
# Select the column containing book titles
documents = data['Book Title'].astype(str)
# Initialize TF-IDF vectors and model
tfidf_doc2vec_model = TFIDFDoc2Vec()
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],
"Edition": data['Edition'].iloc[index],
"File Name": data['File_name'].iloc[index]
}
ranked_documents.append(document_details)
return ranked_documents
# Receive query from the user
#query = input("Enter your query: ")
#result = answer(query)
#print(result)