Use spacy instead of nltk and move data functions to separate module

app/data.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+import bz2
+from typing import Literal
+
+import pandas as pd
+
+from app.constants import (
+    AMAZONREVIEWS_PATH,
+    AMAZONREVIEWS_URL,
+    IMDB50K_PATH,
+    IMDB50K_URL,
+    SENTIMENT140_PATH,
+    SENTIMENT140_URL,
+)
+
+__all__ = ["load_data"]
+
+
+def load_sentiment140(include_neutral: bool = False) -> tuple[list[str], list[int]]:
+    """Load the sentiment140 dataset and make it suitable for use.
+
+    Args:
+        include_neutral: Whether to include neutral sentiment
+
+    Returns:
+        Text and label data
+
+    Raises:
+        FileNotFoundError: If the dataset is not found
+    """
+    # Check if the dataset exists
+    if not SENTIMENT140_PATH.exists():
+        msg = (
+            f"Sentiment140 dataset not found at: '{SENTIMENT140_PATH}'\n"
+            "Please download the dataset from:\n"
+            f"{SENTIMENT140_URL}"
+        )
+        raise FileNotFoundError(msg)
+
+    # Load the dataset
+    data = pd.read_csv(
+        SENTIMENT140_PATH,
+        encoding="ISO-8859-1",
+        names=[
+            "target",  # 0 = negative, 2 = neutral, 4 = positive
+            "id",  # The id of the tweet
+            "date",  # The date of the tweet
+            "flag",  # The query, NO_QUERY if not present
+            "user",  # The user that tweeted
+            "text",  # The text of the tweet
+        ],
+    )
+
+    # Ignore rows with neutral sentiment
+    if not include_neutral:
+        data = data[data["target"] != 2]
+
+    # Map sentiment values
+    data["sentiment"] = data["target"].map(
+        {
+            0: 0,  # Negative
+            4: 1,  # Positive
+            2: 2,  # Neutral
+        },
+    )
+
+    # Return as lists
+    return data["text"].tolist(), data["sentiment"].tolist()
+
+
+def load_amazonreviews(merge: bool = True) -> tuple[list[str], list[int]]:
+    """Load the amazonreviews dataset and make it suitable for use.
+
+    Args:
+        merge: Whether to merge the test and train datasets (otherwise ignore test)
+
+    Returns:
+        Text and label data
+
+    Raises:
+        FileNotFoundError: If the dataset is not found
+    """
+    # Check if the dataset exists
+    test_exists = AMAZONREVIEWS_PATH[0].exists() or not merge
+    train_exists = AMAZONREVIEWS_PATH[1].exists()
+    if not (test_exists and train_exists):
+        msg = (
+            f"Amazonreviews dataset not found at: '{AMAZONREVIEWS_PATH[0]}' and '{AMAZONREVIEWS_PATH[1]}'\n"
+            "Please download the dataset from:\n"
+            f"{AMAZONREVIEWS_URL}"
+        )
+        raise FileNotFoundError(msg)
+
+    # Load the datasets
+    with bz2.BZ2File(AMAZONREVIEWS_PATH[1]) as train_file:
+        train_data = [line.decode("utf-8") for line in train_file]
+
+    test_data = []
+    if merge:
+        with bz2.BZ2File(AMAZONREVIEWS_PATH[0]) as test_file:
+            test_data = [line.decode("utf-8") for line in test_file]
+
+    # Merge the datasets
+    data = train_data + test_data
+
+    # Split the data into labels and text
+    labels, texts = zip(*(line.split(" ", 1) for line in data))
+
+    # Map sentiment values
+    sentiments = [int(label.split("__label__")[1]) - 1 for label in labels]
+
+    # Return as lists
+    return texts, sentiments
+
+
+def load_imdb50k() -> tuple[list[str], list[int]]:
+    """Load the imdb50k dataset and make it suitable for use.
+
+    Returns:
+        Text and label data
+
+    Raises:
+        FileNotFoundError: If the dataset is not found
+    """
+    # Check if the dataset exists
+    if not IMDB50K_PATH.exists():
+        msg = (
+            f"IMDB50K dataset not found at: '{IMDB50K_PATH}'\n"
+            "Please download the dataset from:\n"
+            f"{IMDB50K_URL}"
+        )  # fmt: off
+        raise FileNotFoundError(msg)
+
+    # Load the dataset
+    data = pd.read_csv(IMDB50K_PATH)
+
+    # Map sentiment values
+    data["sentiment"] = data["sentiment"].map(
+        {
+            "positive": 1,
+            "negative": 0,
+        },
+    )
+
+    # Return as lists
+    return data["review"].tolist(), data["sentiment"].tolist()
+
+
+def load_data(dataset: Literal["sentiment140", "amazonreviews", "imdb50k"]) -> tuple[list[str], list[int]]:
+    """Load and preprocess the specified dataset.
+
+    Args:
+        dataset: Dataset to load
+
+    Returns:
+        Text and label data
+
+    Raises:
+        ValueError: If the dataset is not recognized
+    """
+    match dataset:
+        case "sentiment140":
+            return load_sentiment140(include_neutral=False)
+        case "amazonreviews":
+            return load_amazonreviews(merge=True)
+        case "imdb50k":
+            return load_imdb50k()
+        case _:
+            msg = f"Unknown dataset: {dataset}"
+            raise ValueError(msg)

app/model.py

@@ -1,250 +1,86 @@
 from __future__ import annotations
 
-import bz2
-import re
 import warnings
-from typing import Literal
 
-import nltk
-import pandas as pd
+import numpy as np
+import spacy
 from joblib import Memory
-from nltk.corpus import stopwords
-from nltk.stem import WordNetLemmatizer
 from sklearn.base import BaseEstimator, TransformerMixin
-from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
+from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.linear_model import LogisticRegression
-from sklearn.model_selection import cross_val_score, train_test_split
+from sklearn.model_selection import RandomizedSearchCV, cross_val_score, train_test_split
 from sklearn.pipeline import Pipeline
+from tqdm import tqdm
 
-from app.constants import (
-    AMAZONREVIEWS_PATH,
-    AMAZONREVIEWS_URL,
-    CACHE_DIR,
-    EMOTICON_MAP,
-    IMDB50K_PATH,
-    IMDB50K_URL,
-    SENTIMENT140_PATH,
-    SENTIMENT140_URL,
-    URL_REGEX,
-)
+from app.constants import CACHE_DIR
 
-__all__ = ["load_data", "create_model", "train_model", "evaluate_model"]
+__all__ = ["create_model", "train_model", "evaluate_model"]
 
+nlp = spacy.load("en_core_web_sm", disable=["tok2vec", "parser", "ner"])
 
-class TextCleaner(BaseEstimator, TransformerMixin):
+
+class TextTokenizer(BaseEstimator, TransformerMixin):
     def __init__(
         self,
         *,
-        replace_url: bool = True,
-        replace_hashtag: bool = True,
-        replace_emoticon: bool = True,
-        replace_emoji: bool = True,
-        lowercase: bool = True,
         character_threshold: int = 2,
-        remove_special_characters: bool = True,
-        remove_extra_spaces: bool = True,
-    ):
-        self.replace_url = replace_url
-        self.replace_hashtag = replace_hashtag
-        self.replace_emoticon = replace_emoticon
-        self.replace_emoji = replace_emoji
-        self.lowercase = lowercase
+        batch_size: int = 1024,
+        n_jobs: int = 8,
+        progress: bool = True,
+    ) -> None:
         self.character_threshold = character_threshold
-        self.remove_special_characters = remove_special_characters
-        self.remove_extra_spaces = remove_extra_spaces
-
-    def fit(self, _data: list[str], _labels: list[int] | None = None) -> TextCleaner:
-        return self
+        self.batch_size = batch_size
+        self.n_jobs = n_jobs
+        self.progress = progress
 
-    def transform(self, data: list[str], _labels: list[int] | None = None) -> list[str]:
-        # Replace URLs, hashtags, emoticons, and emojis
-        data = [re.sub(URL_REGEX, "URL", text) for text in data] if self.replace_url else data
-        data = [re.sub(r"#\w+", "HASHTAG", text) for text in data] if self.replace_hashtag else data
-
-        # Replace emoticons
-        if self.replace_emoticon:
-            for word, emoticons in EMOTICON_MAP.items():
-                for emoticon in emoticons:
-                    data = [text.replace(emoticon, f"EMOTE_{word}") for text in data]
-
-        # Basic text cleaning
-        data = [text.lower() for text in data] if self.lowercase else data  # Lowercase
-        threshold_pattern = re.compile(rf"\b\w{{1,{self.character_threshold}}}\b")
-        data = (
-            [re.sub(threshold_pattern, "", text) for text in data] if self.character_threshold > 0 else data
-        )  # Remove short words
-        data = (
-            [re.sub(r"[^a-zA-Z0-9\s]", "", text) for text in data] if self.remove_special_characters else data
-        )  # Remove special characters
-        data = [re.sub(r"\s+", " ", text) for text in data] if self.remove_extra_spaces else data  # Remove extra spaces
-
-        # Remove leading and trailing whitespace
-        return [text.strip() for text in data]
-
-
-class TextLemmatizer(BaseEstimator, TransformerMixin):
-    def __init__(self):
-        self.lemmatizer = WordNetLemmatizer()
-
-    def fit(self, _data: list[str], _labels: list[int] | None = None) -> TextLemmatizer:
+    def fit(self, _data: list[str], _labels: list[int] | None = None) -> TextTokenizer:
         return self
 
-    def transform(self, data: list[str], _labels: list[int] | None = None) -> list[str]:
-        return [self.lemmatizer.lemmatize(text) for text in data]
-
-
-def load_sentiment140(include_neutral: bool = False) -> tuple[list[str], list[int]]:
-    """Load the sentiment140 dataset and make it suitable for use.
-
-    Args:
-        include_neutral: Whether to include neutral sentiment
-
-    Returns:
-        Text and label data
-
-    Raises:
-        FileNotFoundError: If the dataset is not found
-    """
-    # Check if the dataset exists
-    if not SENTIMENT140_PATH.exists():
-        msg = (
-            f"Sentiment140 dataset not found at: '{SENTIMENT140_PATH}'\n"
-            "Please download the dataset from:\n"
-            f"{SENTIMENT140_URL}"
-        )
-        raise FileNotFoundError(msg)
-
-    # Load the dataset
-    data = pd.read_csv(
-        SENTIMENT140_PATH,
-        encoding="ISO-8859-1",
-        names=[
-            "target",  # 0 = negative, 2 = neutral, 4 = positive
-            "id",  # The id of the tweet
-            "date",  # The date of the tweet
-            "flag",  # The query, NO_QUERY if not present
-            "user",  # The user that tweeted
-            "text",  # The text of the tweet
-        ],
-    )
-
-    # Ignore rows with neutral sentiment
-    if not include_neutral:
-        data = data[data["target"] != 2]
-
-    # Map sentiment values
-    data["sentiment"] = data["target"].map(
-        {
-            0: 0,  # Negative
-            4: 1,  # Positive
-            2: 2,  # Neutral
-        },
-    )
-
-    # Return as lists
-    return data["text"].tolist(), data["sentiment"].tolist()
-
-
-def load_amazonreviews(merge: bool = True) -> tuple[list[str], list[int]]:
-    """Load the amazonreviews dataset and make it suitable for use.
-
-    Args:
-        merge: Whether to merge the test and train datasets (otherwise ignore test)
-
-    Returns:
-        Text and label data
-
-    Raises:
-        FileNotFoundError: If the dataset is not found
-    """
-    # Check if the dataset exists
-    test_exists = AMAZONREVIEWS_PATH[0].exists() or not merge
-    train_exists = AMAZONREVIEWS_PATH[1].exists()
-    if not (test_exists and train_exists):
-        msg = (
-            f"Amazonreviews dataset not found at: '{AMAZONREVIEWS_PATH[0]}' and '{AMAZONREVIEWS_PATH[1]}'\n"
-            "Please download the dataset from:\n"
-            f"{AMAZONREVIEWS_URL}"
-        )
-        raise FileNotFoundError(msg)
-
-    # Load the datasets
-    with bz2.BZ2File(AMAZONREVIEWS_PATH[1]) as train_file:
-        train_data = [line.decode("utf-8") for line in train_file]
-
-    test_data = []
-    if merge:
-        with bz2.BZ2File(AMAZONREVIEWS_PATH[0]) as test_file:
-            test_data = [line.decode("utf-8") for line in test_file]
+    def transform(self, data: list[str]) -> list[list[str]]:
+        tokenized = []
+        for doc in tqdm(
+            nlp.pipe(data, batch_size=self.batch_size, n_process=self.n_jobs),
+            total=len(data),
+            disable=not self.progress,
+        ):
+            tokens = []
+            for token in doc:
+                # Ignore stop words and punctuation
+                if token.is_stop or token.is_punct:
+                    continue
+                # Ignore emails, URLs and numbers
+                if token.like_email or token.like_email or token.like_num:
+                    continue
 
-    # Merge the datasets
-    data = train_data + test_data
+                # Lemmatize and lowercase
+                tok = token.lemma_.lower().strip()
 
-    # Split the data into labels and text
-    labels, texts = zip(*(line.split(" ", 1) for line in data))
+                # Format hashtags
+                if tok.startswith("#"):
+                    tok = tok[1:]
 
-    # Map sentiment values
-    sentiments = [int(label.split("__label__")[1]) - 1 for label in labels]
+                # Ignore short and non-alphanumeric tokens
+                if len(tok) < self.character_threshold or not tok.isalnum():
+                    continue
 
-    # Return as lists
-    return texts, sentiments
+                # TODO: Emoticons and emojis
+                # TODO: Spelling correction
 
+                tokens.append(tok)
+            tokenized.append(tokens)
+        return tokenized
 
-def load_imdb50k() -> tuple[list[str], list[int]]:
-    """Load the imdb50k dataset and make it suitable for use.
 
-    Returns:
-        Text and label data
-
-    Raises:
-        FileNotFoundError: If the dataset is not found
-    """
-    # Check if the dataset exists
-    if not IMDB50K_PATH.exists():
-        msg = (
-            f"IMDB50K dataset not found at: '{IMDB50K_PATH}'\n"
-            "Please download the dataset from:\n"
-            f"{IMDB50K_URL}"
-        )  # fmt: off
-        raise FileNotFoundError(msg)
-
-    # Load the dataset
-    data = pd.read_csv(IMDB50K_PATH)
-
-    # Map sentiment values
-    data["sentiment"] = data["sentiment"].map(
-        {
-            "positive": 1,
-            "negative": 0,
-        },
-    )
-
-    # Return as lists
-    return data["review"].tolist(), data["sentiment"].tolist()
-
-
-def load_data(dataset: Literal["sentiment140", "amazonreviews", "imdb50k"]) -> tuple[list[str], list[int]]:
-    """Load and preprocess the specified dataset.
+def identity(x: list[str]) -> list[str]:
+    """Identity function for use in TfidfVectorizer.
 
     Args:
-        dataset: Dataset to load
+        x: Input data
 
     Returns:
-        Text and label data
-
-    Raises:
-        ValueError: If the dataset is not recognized
+        Unchanged input data
     """
-    match dataset:
-        case "sentiment140":
-            return load_sentiment140(include_neutral=False)
-        case "amazonreviews":
-            return load_amazonreviews(merge=True)
-        case "imdb50k":
-            return load_imdb50k()
-        case _:
-            msg = f"Unknown dataset: {dataset}"
-            raise ValueError(msg)
+    return x
 
 
 def create_model(
@@ -262,26 +98,22 @@ def create_model(
     Returns:
         Untrained model
     """
-    # Download NLTK data if not already downloaded
-    nltk.download("wordnet", quiet=True)
-    nltk.download("stopwords", quiet=True)
-
-    # Load English stopwords
-    stopwords_en = set(stopwords.words("english"))
-
     return Pipeline(
         [
-            # Text preprocessing
-            ("clean", TextCleaner()),
-            ("lemma", TextLemmatizer()),
-            # Preprocess (NOTE: Can be replaced with TfidfVectorizer, but left for clarity)
+            ("tokenizer", TextTokenizer(progress=True)),
             (
-                "vectorize",
-                CountVectorizer(stop_words=stopwords_en, ngram_range=(1, 2), max_features=max_features),
+                "vectorizer",
+                TfidfVectorizer(
+                    max_features=max_features,
+                    ngram_range=(1, 2),
+                    # disable text processing
+                    tokenizer=identity,
+                    preprocessor=identity,
+                    lowercase=False,
+                    token_pattern=None,
+                ),
             ),
-            ("tfidf", TfidfTransformer()),
-            # Classifier
-            ("clf", LogisticRegression(max_iter=1000, random_state=seed)),
+            ("classifier", LogisticRegression(max_iter=1000, C=1.0, random_state=seed)),
         ],
         memory=Memory(CACHE_DIR, verbose=0),
         verbose=verbose,
@@ -289,11 +121,11 @@ def create_model(
 
 
 def train_model(
-    model: Pipeline,
+    model: BaseEstimator,
     text_data: list[str],
     label_data: list[int],
     seed: int = 42,
-) -> tuple[float, list[str], list[int]]:
+) -> tuple[BaseEstimator, float]:
     """Train the sentiment analysis model.
 
     Args:
@@ -303,7 +135,7 @@ def train_model(
         seed: Random seed (None for random seed)
 
     Returns:
-        Model accuracy and test data
+        Trained model and accuracy
     """
     text_train, text_test, label_train, label_test = train_test_split(
         text_data,
@@ -312,37 +144,52 @@ def train_model(
         random_state=seed,
     )
 
+    param_distributions = {
+        "classifier__C": np.logspace(-4, 4, 20),
+        "classifier__penalty": ["l1", "l2"],
+    }
+
+    search = RandomizedSearchCV(
+        model,
+        param_distributions,
+        n_iter=10,
+        cv=5,
+        scoring="accuracy",
+        random_state=seed,
+        n_jobs=-1,
+    )
+
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
-        model.fit(text_train, label_train)
+        # model.fit(text_train, label_train)
+        search.fit(text_train, label_train)
 
-    return model.score(text_test, label_test), text_test, label_test
+    best_model = search.best_estimator_
+    return best_model, best_model.score(text_test, label_test)
 
 
 def evaluate_model(
     model: Pipeline,
-    text_test: list[str],
-    label_test: list[int],
-    cv: int = 5,
+    text_data: list[str],
+    label_data: list[int],
+    folds: int = 5,
 ) -> tuple[float, float]:
     """Evaluate the model using cross-validation.
 
     Args:
         model: Trained model
-        text_test: Text data
-        label_test: Label data
-        seed: Random seed (None for random seed)
-        cv: Number of cross-validation folds
+        text_data: Text data
+        label_data: Label data
+        folds: Number of cross-validation folds
 
     Returns:
         Mean accuracy and standard deviation
     """
     scores = cross_val_score(
         model,
-        text_test,
-        label_test,
-        cv=cv,
+        text_data,
+        label_data,
+        cv=folds,
         scoring="accuracy",
-        n_jobs=-1,
     )
     return scores.mean(), scores.std()