From b9b749cdd4696aac264f88d3eeb5b74a261ba586 Mon Sep 17 00:00:00 2001 From: "Akeem T. L. King" Date: Tue, 28 Jul 2026 17:43:00 -0400 Subject: [PATCH 1/2] Make NLP classification cheaper during crawls --- src/torbot/modules/nlp/README.md | 35 ++++-- src/torbot/modules/nlp/gather_data.py | 38 ++++--- src/torbot/modules/nlp/main.py | 149 +++++++++++++++++++------- tests/test_nlp.py | 77 +++++++++++++ 4 files changed, 237 insertions(+), 62 deletions(-) create mode 100644 tests/test_nlp.py diff --git a/src/torbot/modules/nlp/README.md b/src/torbot/modules/nlp/README.md index 0073143c..5697c220 100644 --- 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("...") +``` + +`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 index f47750e5..f20cb3f9 100644 --- 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 index 9722c601..1c04f088 100644 --- 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 new file mode 100644 index 00000000..489e2078 --- /dev/null +++ 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("

Hello

World

") == "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") From d4c5b1c22f9de040d8139956d3651a5772ffd616 Mon Sep 17 00:00:00 2001 From: "Akeem T. L. King" Date: Tue, 28 Jul 2026 17:43:34 -0400 Subject: [PATCH 2/2] Bump version for the NLP cleanup --- docs/CHANGELOG.md | 13 +++++++++++++ pyproject.toml | 2 +- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index b995c9d3..353936d7 100644 --- 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 index 264b1ac7..6f8e52c2 100644 --- 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" },