DedSecInside · KingAkeem · Jul 28, 2026 · Jul 28, 2026 · Jul 28, 2026
diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
@@ -2,6 +2,19 @@
 --------------------
 All notable changes to this project will be documented in this file.
 
+## 4.3.0 - 2026-07-28
+
+### Changed
+- Made NLP classification train lazily from the packaged CSV and cache the classifier for the process, avoiding model retraining for every crawled page.
+- Changed the NLP score returned by `classify()` from a misleading random-split accuracy value to prediction confidence for the selected category.
+
+### Fixed
+- Removed runtime `training_data/` generation from page classification so installed packages do not try to write generated files into `site-packages`.
+- Removed process-wide directory changes from the NLP helpers.
+
+### Tests
+- Added direct NLP coverage for HTML text extraction, CSV loading, classifier scoring, and path-safe training-data export.
+
 ## 4.2.0 - 2026-07-28
 
 Install from PyPI after publishing with:

diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "torbot"
-version = "4.2.0"
+version = "4.3.0"
 authors = [
     { name="Akeem King", email="akeemtlking@gmail.com" },
     { name="PS Narayanan", email="thepsnarayanan@gmail.com" },

diff --git a/src/torbot/modules/nlp/README.md b/src/torbot/modules/nlp/README.md
@@ -1,11 +1,30 @@
-# Natural Language Processing Library
+# Natural Language Processing Module
 
-This library provides tool for performing natural language processing on websites.
-This library is in it's infancy currently and can only be used for testing.
+This module classifies crawled pages into broad website categories using the
+packaged `website_classification.csv` dataset.
 
-To test gathering data use:
-`python3 gater_data.py`
-* This will generate the data necessary to train the classification model 
+Runtime classification uses `classify(html)` from `main.py`. The classifier is
+trained lazily from the CSV on first use and cached for the rest of the process,
+so crawling multiple pages does not retrain the model for every page.
 
-To predict the classification of a webiste use:
-`classify.py` and provide the url 
+## Public API
+
+```python
+from torbot.modules.nlp.main import classify
+
+category, confidence = classify("<html>...</html>")
+```
+
+`confidence` is the model's confidence for the selected category. It is not a
+model-wide accuracy score.
+
+## Optional training-data export
+
+The crawler no longer needs a generated `training_data/` directory. If you need
+one for experiments with `sklearn.datasets.load_files`, run:
+
+```sh
+python3 -m torbot.modules.nlp.gather_data
+```
+
+This creates an ignored `training_data/` directory beside the NLP module.
diff --git a/src/torbot/modules/nlp/gather_data.py b/src/torbot/modules/nlp/gather_data.py
@@ -1,15 +1,19 @@
 import csv
-import os
-
 from pathlib import Path
+from typing import Union
+
 
-os.chdir(Path(__file__).parent)
+NLP_DIRECTORY = Path(__file__).resolve().parent
+DEFAULT_CSV_PATH = NLP_DIRECTORY / "website_classification.csv"
+DEFAULT_OUTPUT_DIRECTORY = NLP_DIRECTORY / "training_data"
 
 
-def write_data():
+def write_data(
+    csv_path: Union[str, Path] = DEFAULT_CSV_PATH,
+    output_directory: Union[str, Path] = DEFAULT_OUTPUT_DIRECTORY,
+) -> None:
     """
-    Writes the training data from the csv file to a directory based on the
-    scikit-learn.datasets `load_files` specification.
+    Write CSV rows to a scikit-learn load_files-compatible directory tree.
 
     dataset source: https://www.kaggle.com/hetulmehta/website-classification
 
@@ -20,17 +24,19 @@ def write_data():
             category_2_folder/
                     file_43.txt file_44.txt ...
     """
-
-    with open("website_classification.csv") as csvfile:
-        website_reader = csv.reader(csvfile, delimiter=",")
+    output_path = Path(output_directory)
+    with Path(csv_path).open(newline="", encoding="utf-8", errors="replace") as csvfile:
+        website_reader = csv.DictReader(csvfile)
         for row in website_reader:
-            [id, website, content, category] = row
-            if category != "category":
-                category = category.replace("/", "+")
-            dir_name = f"training_data/{category}"
-            Path(dir_name).mkdir(parents=True, exist_ok=True)
-            with open(f"{dir_name}/{id}.txt", mode="w+") as txtfile:
-                txtfile.write(content)
+            row_id = (row.get("id") or "").strip()
+            content = row.get("cleaned_text") or ""
+            category = (row.get("category") or "").strip().replace("/", "+")
+            if not row_id or not content or not category:
+                continue
+
+            category_directory = output_path / category
+            category_directory.mkdir(parents=True, exist_ok=True)
+            (category_directory / f"{row_id}.txt").write_text(content, encoding="utf-8")
 
 
 if __name__ == "__main__":

diff --git a/src/torbot/modules/nlp/main.py b/src/torbot/modules/nlp/main.py
@@ -1,48 +1,121 @@
-import numpy as np
-import os
+"""
+Website text classification helpers.
+
+The classifier is trained lazily from the packaged CSV and cached for the
+process lifetime. That keeps LinkTree classification cheap after the first
+call and avoids writing generated training files into the installed package.
+"""
+from __future__ import annotations
+
+import csv
+from functools import lru_cache
 from pathlib import Path
+from typing import List, Tuple, Union
 
+import numpy as np
 from bs4 import BeautifulSoup
-from sklearn.model_selection import train_test_split
-from sklearn.pipeline import Pipeline
+from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.linear_model import SGDClassifier
-from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
-from sklearn.datasets import load_files
+from sklearn.pipeline import Pipeline
 
 
-def classify(data):
-    """
-    Classify URL specified by user
-    """
-    soup = BeautifulSoup(data, features="html.parser")
-    html = soup.get_text()
+NLP_DIRECTORY = Path(__file__).resolve().parent
+DEFAULT_DATASET_PATH = NLP_DIRECTORY / "website_classification.csv"
+
+
+def extract_text(html: str) -> str:
+    """Extract normalized visible text from an HTML document."""
+    soup = BeautifulSoup(html or "", features="html.parser")
+    return soup.get_text(" ", strip=True)
+
+
+def load_training_rows(
+    csv_path: Union[str, Path] = DEFAULT_DATASET_PATH,
+) -> Tuple[List[str], List[str]]:
+    """Load training text and labels from the website classification CSV."""
+    path = Path(csv_path)
+    texts: list[str] = []
+    labels: list[str] = []
+
+    with path.open(newline="", encoding="utf-8", errors="replace") as csvfile:
+        reader = csv.DictReader(csvfile)
+        for row in reader:
+            text = (row.get("cleaned_text") or "").strip()
+            category = (row.get("category") or "").strip()
+            if text and category:
+                texts.append(text)
+                labels.append(category)
+
+    if not texts:
+        raise ValueError(f"no training rows found in {path}")
+
+    if len(set(labels)) < 2:
+        raise ValueError(f"at least two categories are required in {path}")
 
-    # create classifier
-    clf = Pipeline(
+    return texts, labels
+
+
+def build_classifier(csv_path: Union[str, Path] = DEFAULT_DATASET_PATH) -> Pipeline:
+    """Train a text classifier from the supplied dataset CSV."""
+    texts, labels = load_training_rows(csv_path)
+    classifier = Pipeline(
         [
-            ("vect", CountVectorizer()),
-            ("tfidf", TfidfTransformer()),
-            ("clf", SGDClassifier()),
+            (
+                "tfidf",
+                TfidfVectorizer(
+                    lowercase=True,
+                    max_features=50000,
+                    ngram_range=(1, 2),
+                    sublinear_tf=True,
+                ),
+            ),
+            (
+                "clf",
+                SGDClassifier(
+                    class_weight="balanced",
+                    loss="modified_huber",
+                    max_iter=1000,
+                    random_state=42,
+                    tol=1e-3,
+                ),
+            ),
         ]
     )
-    try:
-        os.chdir(Path(__file__).parent)
-
-        dataset = load_files("training_data")
-    except FileNotFoundError:
-        print("Training data not found. Obtaining training data...")
-        print("This may take a while...")
-        from .gather_data import write_data
-
-        write_data()
-        print("Training data obtained.")
-        dataset = load_files("training_data")
-        pass
-    x_train, x_test, y_train, y_test = train_test_split(dataset.data, dataset.target)
-    clf.fit(x_train, y_train)
-
-    # returns an array of target_name values
-    predicted = clf.predict([html])
-    accuracy = np.mean(predicted == y_test)
-
-    return [dataset.target_names[predicted[0]], accuracy]
+    classifier.fit(texts, labels)
+    return classifier
+
+
+@lru_cache(maxsize=4)
+def get_classifier(csv_path: Union[str, Path] = DEFAULT_DATASET_PATH) -> Pipeline:
+    """Return a cached classifier for the requested dataset."""
+    return build_classifier(Path(csv_path).resolve())
+
+
+def _prediction_score(classifier: Pipeline, text: str) -> float:
+    """Return a best-effort confidence score for a single prediction."""
+    if hasattr(classifier, "predict_proba"):
+        probabilities = classifier.predict_proba([text])[0]
+        return float(np.max(probabilities))
+
+    if hasattr(classifier, "decision_function"):
+        scores = np.asarray(classifier.decision_function([text]))
+        return float(np.max(scores))
+
+    return 0.0
+
+
+def classify(data: str) -> List[Union[str, float]]:
+    """
+    Classify the supplied HTML and return [category, confidence].
+
+    The confidence value is model confidence for the selected category, not a
+    model-wide accuracy score.
+    """
+    text = extract_text(data)
+    if not text:
+        return ["unknown", 0.0]
+
+    classifier = get_classifier()
+    prediction = classifier.predict([text])[0]
+    confidence = _prediction_score(classifier, text)
+    return [str(prediction), confidence]
diff --git a/tests/test_nlp.py b/tests/test_nlp.py
@@ -0,0 +1,77 @@
+import importlib
+import os
+import sys
+
+
+def import_real_nlp_module():
+    """Import the real NLP module instead of the lightweight test stub."""
+    sys.modules.pop("torbot.modules.nlp.main", None)
+    return importlib.import_module("torbot.modules.nlp.main")
+
+
+def write_fixture_csv(path):
+    path.write_text(
+        "\n".join(
+            [
+                "id,website,cleaned_text,category",
+                "1,https://sports.test,basketball team score playoff coach,Sports",
+                "2,https://sports2.test,football league match tournament goal,Sports",
+                "3,https://biz.test,revenue company shareholder market,Business",
+                "4,https://biz2.test,corporate earnings sales customer,Business",
+            ]
+        ),
+        encoding="utf-8",
+    )
+
+
+def test_extract_text_normalizes_html() -> None:
+    nlp = import_real_nlp_module()
+
+    assert nlp.extract_text("<html><h1>Hello</h1><p>World</p></html>") == "Hello World"
+
+
+def test_load_training_rows_uses_csv_without_training_directory(tmp_path) -> None:
+    nlp = import_real_nlp_module()
+    csv_path = tmp_path / "website_classification.csv"
+    write_fixture_csv(csv_path)
+
+    texts, labels = nlp.load_training_rows(csv_path)
+
+    assert len(texts) == 4
+    assert sorted(set(labels)) == ["Business", "Sports"]
+    assert not (tmp_path / "training_data").exists()
+
+
+def test_build_classifier_returns_category_and_confidence(tmp_path) -> None:
+    nlp = import_real_nlp_module()
+    csv_path = tmp_path / "website_classification.csv"
+    write_fixture_csv(csv_path)
+
+    classifier = nlp.build_classifier(csv_path)
+    text = "basketball team wins playoff tournament"
+    prediction = classifier.predict([text])[0]
+    confidence = nlp._prediction_score(classifier, text)
+
+    assert prediction == "Sports"
+    assert 0.0 <= confidence <= 1.0
+
+
+def test_classify_empty_html_returns_unknown() -> None:
+    nlp = import_real_nlp_module()
+
+    assert nlp.classify("") == ["unknown", 0.0]
+
+
+def test_gather_data_writes_training_files_without_changing_cwd(tmp_path) -> None:
+    from torbot.modules.nlp.gather_data import write_data
+
+    csv_path = tmp_path / "website_classification.csv"
+    output_path = tmp_path / "training_data"
+    write_fixture_csv(csv_path)
+    before = os.getcwd()
+
+    write_data(csv_path=csv_path, output_directory=output_path)
+
+    assert os.getcwd() == before
+    assert (output_path / "Sports" / "1.txt").read_text(encoding="utf-8")
+    assert (output_path / "Business" / "3.txt").read_text(encoding="utf-8")