preprocess.py

# Main imports.
import string
import json
import math
from pathlib import Path
from typing import Dict, List

import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from nltk.stem.porter import PorterStemmer

# Ensure required NLTK data is available
nltk.download('stopwords', quiet=True)
nltk.download('punkt', quiet=True)
nltk.download('punkt_tab', quiet=True)


# ----------------------------------------------------------------------
# NLTK setup
# ----------------------------------------------------------------------

def _ensure_nltk_resources() -> None:
    """Ensure required NLTK data is available, downloading quietly if needed."""
    resources = {
        "stopwords": "corpora/stopwords",
        "punkt": "tokenizers/punkt",
    }

    for res, path in resources.items():
        try:
            nltk.data.find(path)
        except LookupError:
            nltk.download(res, quiet=True)


_ensure_nltk_resources()

# ----------------------------------------------------------------------
# Paths & globals
# ----------------------------------------------------------------------

ASSETS_DIR = Path("./assets")
TWEET_FILE = ASSETS_DIR / "tweet_list.txt"
QUERY_FILE = ASSETS_DIR / "test_queries.txt"
STOPWORDS_FILE = ASSETS_DIR / "stop_words.txt"

# Initialize stemmer
ps = PorterStemmer()

# Custom Stopwords that are NOT defined in Library or Provided Stopwords
EDGE_STOPWORDS = {"n't", "'d", "http", "https", "//", "..."}


def _load_stopwords() -> set:
    """Load and cache stopwords (library + custom file + edge cases)."""
    words = set(stopwords.words("english"))

    if STOPWORDS_FILE.is_file():
        with STOPWORDS_FILE.open(encoding="utf-8") as f:
            for line in f:
                w = line.strip()
                if w:
                    words.add(w)

    words.update(EDGE_STOPWORDS)
    return words


CUSTOM_STOPWORDS = _load_stopwords()

# ----------------------------------------------------------------------
# Utilities
# ----------------------------------------------------------------------

def isNumeric(subj: str) -> bool:
    """
    Check if a string contains numerical values.

    :param str subj: the string to be tested
    :return: True if subj is numeric, False otherwise.
    :rtype: bool
    """
    try:
        float(subj)
        return True
    except (ValueError, TypeError):
        return False


def filterSentence(sentence: str, verbose: bool = False) -> List[str]:
    """
    Step 1: Filters sentences from tweets and queries.

    TODO: Filter URLs, non-English characters, punctuation embedded in words,
    and unicode normalization.

    :param str sentence: The sentence to be tokenized with stopwords and punctuation removed.
    :param bool verbose: [Optional] Provide printed output of tokens for testing.
    :return: the tokenized sentence.
    :rtype: list[str]
    """
    tokens: List[str] = []

    for word in word_tokenize(sentence):
        lower = word.lower()

        if (
            lower not in CUSTOM_STOPWORDS
            and word not in string.punctuation
            and not isNumeric(lower)
        ):
            tokens.append(ps.stem(lower))

    if verbose:
        print("\nTesting string:\n\n\t", sentence, "\n")
        print("Tokenized:\n")
        print("\t[{}]\n".format(", ".join(tokens)))

    return tokens


# ----------------------------------------------------------------------
# Input / import helpers
# ----------------------------------------------------------------------

def importTweets(verbose: bool = False) -> Dict[str, List[str]]:
    """
    Import tweets from collection.

    :param bool verbose: [Optional] Provide printed output of tokens for testing.
    :return: mapping from tweet_id -> tokenized tweet.
    :rtype: dict[str, list[str]]
    """
    tweet_list: Dict[str, List[str]] = {}

    if not TWEET_FILE.is_file():
        raise FileNotFoundError(f"Tweet file not found: {TWEET_FILE}")

    # Splits tweet list at newline character.
    with TWEET_FILE.open("r", encoding="utf-8-sig") as f:
        for line in f:
            line = line.rstrip("\n")
            if not line:
                continue

            # Expect: "<id>\t<text>"
            try:
                key, value = line.split("\t", 1)
            except ValueError:
                # Malformed line, skip it instead of blowing up
                continue

            tweet_list[key] = filterSentence(value, verbose)

    return tweet_list


def importQuery(verbose: bool = False) -> Dict[int, List[str]]:
    """
    Import query from collection.

    Assumes each query block is separated by a blank line and contains:
      <title> ... </title>

    :param bool verbose: [Optional] Provide printed output of tokens for testing.
    :return: mapping from query_id -> tokenized title.
    :rtype: dict[int, list[str]]
    """
    query_list: Dict[int, List[str]] = {}

    if not QUERY_FILE.is_file():
        raise FileNotFoundError(f"Query file not found: {QUERY_FILE}")

    with QUERY_FILE.open("r", encoding="utf-8") as file:
        file_contents = file.read()

    blocks = file_contents.strip().split("\n\n")

    for query_id, block in enumerate(blocks, start=1):
        start = block.find("<title>")
        end = block.find("</title>", start + len("<title>"))

        if start == -1 or end == -1:
            # No title tags found in this block; skip it.
            continue

        # Extract inner text of <title>...</title>
        raw_title = block[start + len("<title>"):end].strip()

        query_list[query_id] = filterSentence(raw_title, verbose)

    return query_list


# ----------------------------------------------------------------------
# Indexing
# ----------------------------------------------------------------------

def buildIndex(documents: Dict, verbose: bool = False) -> Dict[str, Dict]:
    """
    Step 2: Build an inverted index with TF-IDF weights.

    :param dict documents: mapping from document_id -> list of tokens
    :param bool verbose: [Optional] Provide printed output of tokens for testing.
    :return: An inverted index (token -> {doc_id: tf-idf})
    :rtype: dict[str, dict]
    """
    # term -> {doc_id -> term_frequency}
    term_frequencies: Dict[str, Dict] = {}

    for doc_id, tokens in documents.items():
        for token in tokens:
            doc_freqs = term_frequencies.setdefault(token, {})
            doc_freqs[doc_id] = doc_freqs.get(doc_id, 0) + 1

    num_docs = len(documents)
    inverted_index: Dict[str, Dict] = {}

    # Standard IDF: log(N / df)
    for token, doc_freqs in term_frequencies.items():
        df = len(doc_freqs)
        if df == 0:
            continue

        idf = math.log(num_docs / df, 2)
        # Round idf, then multiply by tf (similar to your original intent)
        idf_rounded = round(idf, 3)

        inverted_index[token] = {
            doc_id: tf * idf_rounded for doc_id, tf in doc_freqs.items()
        }

    if verbose:
        print("\nInverted Index (tf-idf):")
        print(json.dumps(inverted_index, indent=2))
        print("-" * 40)

    return inverted_index


def lengthOfDocument(
    inverted_index: Dict[str, Dict],
    tweets: Dict,
    verbose: bool = False,
) -> Dict:
    """
    Compute the Euclidean length of each document vector.

    :param dict inverted_index: token -> {doc_id: tf-idf}
    :param dict tweets: mapping from doc_id -> list of tokens (for doc IDs)
    :param bool verbose: [Optional] print the document lengths
    :return: mapping from doc_id -> document length
    :rtype: dict
    """
    document_lengths: Dict = {}

    for tweet_id, tokens in tweets.items():
        length_sq = 0.0

        # Use unique tokens; each term contributes once to the vector length
        for token in set(tokens):
            weight = inverted_index.get(token, {}).get(tweet_id)
            if weight:
                length_sq += weight * weight

        document_lengths[tweet_id] = round(math.sqrt(length_sq), 3)

    if verbose:
        print("Length of documents", document_lengths)

    return document_lengths