src/models/gensim_vectorizer.py

import pandas as pd
from gensim import corpora
from gensim import similarities
from gensim.models import TfidfModel
from gensim.parsing import strip_tags, strip_numeric, \
    strip_multiple_whitespaces, stem_text, strip_punctuation, \
    remove_stopwords, preprocess_string
import re

from typing import List
from utils.constants import TEST_INPUTS
import argparse
from random import choice

import sys



SAMPLES = 3000
CORPUS_DICTIONARY_PATH="30Ktokens"
ARXIV_DATASR_PATH = "/Users/luis.morales/Desktop/arxiv-paper-recommender/data/processed/reduced_arxiv_papers.parquet.gzip"
SAVE_DICT = False
QUERY = ""

transform_to_lower = lambda s: s.lower()
remove_single_char = lambda s: re.sub(r'\s+\w{1}\s+', '', s)

cleaning_filters = [
    strip_tags,
    strip_numeric,
    strip_punctuation, 
    strip_multiple_whitespaces, 
    transform_to_lower,
    remove_stopwords,
    remove_single_char
]

def gensim_tokenizer(docs: List[str]):
    """
    Tokenizes a list of strings using a series of cleaning filters.

    Args:
        docs (List[str]): A list of strings to be tokenized.

    Returns:
        List[List[str]]: A list of tokenized documents, where each document is represented as a list of tokens.
    """
    tokenized_docs = list()
    for doc in docs:
        processed_words = preprocess_string(doc, cleaning_filters)
        tokenized_docs.append(processed_words)
    
    return tokenized_docs


def cleaning_pipe(document):
    """
    Applies a series of cleaning steps to a document.

    Args:
        document (str): The document to be cleaned.

    Returns:
        list: A list of processed words after applying the cleaning filters.
    """
    # Invoking gensim.parsing.preprocess_string method with set of filters
    processed_words = preprocess_string(document, cleaning_filters)
    return processed_words


def get_gensim_dictionary(tokenized_docs: List[str], dict_name: str = "corpus", save_dict: bool = False):
    """
        Create dictionary of words in preprocessed corpus and saves the dict object
    """
    dictionary = corpora.Dictionary(tokenized_docs)
    if save_dict:    
        parent_folder = "/Users/luis.morales/Desktop/arxiv-paper-recommender/models/nlp_dictionaries"
        dictionary.save(f'{parent_folder}/{dict_name}.dict')
    return dictionary


def get_closest_n(query: str, n: int):
    '''
    Retrieves the top matching documents as per cosine similarity
    between the TF-IDF vector of the query and all documents.

    Args:
        query (str): The query string to find matching documents.
        n (int): The number of closest documents to retrieve.

    Returns:
        numpy.ndarray: An array of indices representing the top matching documents.
    '''
    # Clean the query document using cleaning_pipe function
    query_document = cleaning_pipe(query)

    # Convert the query document to bag-of-words representation
    query_bow = dictionary.doc2bow(query_document)

    # Calculate similarity scores between the query and all documents using TF-IDF model
    sims = index[tfidf_model[query_bow]]

    # Get the indices of the top n closest documents based on similarity scores
    top_idx = sims.argsort()[-1 * n:][::-1]

    return top_idx


def get_recomendations_metadata(query: str, df: pd.DataFrame, n: int):
    '''
    Retrieves metadata recommendations based on a query using cosine similarity.

    Args:
        query (str): The query string for which recommendations are sought.
        n (int): The number of recommendations to retrieve.
        df (pd.DataFrame): The DataFrame containing metadata information.

    Returns:
        pd.DataFrame: A DataFrame containing the recommended metadata, reset with a new index.
    '''
    # Get the indices of the closest matching documents based on the query
    recommendations_idxs = get_closest_n(query, n)
    
    # Retrieve the recommended metadata rows from the DataFrame based on the indices
    recommendations_metadata = df.iloc[recommendations_idxs]
    
    # Reset the index of the recommended metadata DataFrame
    recommendations_metadata = recommendations_metadata.reset_index(drop=True)
    
    return recommendations_metadata

if __name__ == "__main__":
    """
        Example: 
        python script.py --samples 3000 --corpus_dictionary_path "30Ktokens.dict" --arxiv_datasr_path "/Users/luis.morales/Desktop/arxiv-paper-recommender/data/processed/reduced_arxiv_papers.parquet.gzip" --save_dict --query "your query here"

    """
    # Define and parse command-line arguments
    parser = argparse.ArgumentParser(description='ArXiv Paper Recommender CLI')
    parser.add_argument('--samples', default=30000, type=int, help='Number of samples to consider')
    parser.add_argument('--corpus_dictionary_path', default=None ,type=str, help='Path to the corpus dictionary')
    parser.add_argument('--save_dict', default=False, help='Flag to save the dictionary')
    parser.add_argument('--arxiv_dataset_path',
                        default="/Users/luis.morales/Desktop/arxiv-paper-recommender/data/processed/reduced_arxiv_papers.parquet.gzip",
                        type=str, help='Path to the ARXIV parquet source')
    parser.add_argument('--query', default=None, type=str, help='User query')
    args = parser.parse_args()

    num_samples = args.samples
    corpus_dictionary_path = args.corpus_dictionary_path
    arxiv_dataset_path = args.arxiv_dataset_path
    save_dict = args.save_dict
    query = args.query

    print("Parameters:")
    print(f"num_samples: {num_samples}, type: {type(num_samples)}")
    print(f"corpus_dictionary_path: {corpus_dictionary_path}, type: {type(corpus_dictionary_path)}")
    print(f"arxiv_dataset_path: {arxiv_dataset_path}, type: {type(arxiv_dataset_path)}")
    print(f"save_dict: {save_dict}, type: {type(save_dict)}")
    print(f"query: {query}, type: {type(query)}")
    

    if num_samples is None:
        df = pd.read_parquet(arxiv_dataset_path)
    df = pd.read_parquet(arxiv_dataset_path).sample(num_samples).reset_index(drop=True)
    

    corpus = df['cleaned_abstracts'].to_list()
    tokenized_corpus = gensim_tokenizer(corpus)

    dictionary = get_gensim_dictionary(
        tokenized_docs=tokenized_corpus, 
        dict_name=corpus_dictionary_path, 
        save_dict=save_dict
        )

    BoW_corpus = [dictionary.doc2bow(doc, allow_update=True) for doc in tokenized_corpus]

    tfidf_model = TfidfModel(BoW_corpus)

    index = similarities.SparseMatrixSimilarity(tfidf_model[BoW_corpus], num_features=len(dictionary))

    if query is None:
        query = choice(TEST_INPUTS)

    results_df = get_recomendations_metadata(query=query, df=df, n=3)


    for abstract in list(zip(results_df['abstract'].to_list(), results_df['title'].to_list())):
        print(f"User Request ---- : \n {query}")
        print(f"User Request ---- : \n ")
        print(f"Title: {abstract[1]}")
        print(f"Abstract: {abstract[0]}\n")
        print(f"--------------------------")