ENH Improve `plot_semi_supervised_newsgroups.py` example
@@ -11,7 +11,6 @@ loader or setting them to `None` to get all 20 of them.
"""
-import os
import numpy as np
@@ -26,7 +25,17 @@ from sklearn.semi_supervised import SelfTrainingClassifier
from sklearn.semi_supervised import LabelSpreading
from sklearn.metrics import f1_score
-data = fetch_20newsgroups(subset="train", categories=None)
+# Loading dataset containing first five categories
+data = fetch_20newsgroups(
+ subset="train",
+ categories=[
+ "alt.atheism",
+ "comp.graphics",
+ "comp.os.ms-windows.misc",
+ "comp.sys.ibm.pc.hardware",
+ "comp.sys.mac.hardware",
+ ],
+)
print("%d documents" % len(data.filenames))
print("%d categories" % len(data.target_names))
print()
@@ -98,7 +107,5 @@ if __name__ == "__main__":
print("SelfTrainingClassifier on 20% of the training data (rest is unlabeled):")
eval_and_print_metrics(st_pipeline, X_train, y_train, X_test, y_test)
- if "CI" not in os.environ:
- # LabelSpreading takes too long to run in the online documentation
print("LabelSpreading on 20% of the data (rest is unlabeled):")
eval_and_print_metrics(ls_pipeline, X_train, y_train, X_test, y_test)
