classifier_svm.py

#!/usr/bin/env python

import sys
from functools import partial

import spacy
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import Pipeline
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.multiclass import OneVsRestClassifier
from sklearn import svm
from sklearn.calibration import CalibratedClassifierCV

from svm_utils import (
    load_data,
    print_division,
    print_best_features,
    plot_confusion_matrix,
    parse_args_svm,
)


def main():
    args = parse_args_svm()
    # Nice for reproducibility when running lots of exps
    print("Generated by:\npython {0}\n".format(" ".join(sys.argv)))

    # Read in training data
    X_train, Y_train = load_data(
        args.root_dir, "train", args.use_google_data, args.use_normalized_data
    )
    X_test, Y_test = load_data(
        args.root_dir, "dev", args.use_google_data, args.use_normalized_data
    )

    def _tokenizer(doc, spacy_tokenizer):
        return [str(tkn) for tkn in spacy_tokenizer(doc)]

    # Convert the texts to vectors.
    if args.tfidf:
        vec = TfidfVectorizer(min_df=args.min_df, ngram_range=(1, 2))
    else:
        # Simple BoW vectorizer.
        tokenizer = partial(_tokenizer, spacy_tokenizer=spacy.load("en_core_web_sm"))
        vec = CountVectorizer(
            min_df=args.min_df, ngram_range=(1, 2), tokenizer=tokenizer
        )

    # Choose the algorithm.
    if args.algorithm == "nb":
        clf = Pipeline([("vec", vec), ("cls", MultinomialNB())])
    elif args.algorithm == "svm":
        clf = svm.LinearSVC(C=1)
        # Use the CalibratedClassifier so we can use predict_proba later.
        if args.probabilities:
            clf = CalibratedClassifierCV(base_estimator=clf)
        clf = Pipeline([("vec", vec), ("cls", clf)])

    # Do we do 1v1 or 1 v rest?
    if args.one_vs_rest:
        clf = Pipeline([("cls", OneVsRestClassifier(clf))])

    # Train & test on separate set
    clf.fit(X_train, Y_train)
    Y_pred = clf.predict(X_test)
    print_division(clf.named_steps["cls"].classes_, Y_train)
    print(classification_report(Y_test, Y_pred, digits=3))

    # Feature analysis, only possible for SVM.
    if args.features and args.algorithm == "svm" and not args.one_vs_rest:
        print_best_features(vec, clf, X_train, args.only_word_features)

    # Confusion matrix if we want.
    if args.confusion:
        Y_plot = Y_test
        plot_confusion_matrix(
            confusion_matrix(Y_plot, Y_pred),
            [c for c in clf.named_steps["cls"].classes_],
            args.confusion,
            normalize=False,
        )


if __name__ == "__main__":
    main()