Spaces:
Sleeping
Sleeping
#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], | |
"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 | |
# Receive query from the user | |
#query = input("Enter your query: ") | |
#result = answer(query) | |
#print(result) |