if len(sys.argv) > 1 and sys.argv[1] != "":
        source = sys.argv[1]
    else:
        source = IMDB_DATA_DEFAULT

    print "Loading data from original source"
    imdb = ImdbDataHandler(source=source)
    (train_reviews, train_labels) = imdb.get_data(type=ImdbDataHandler.DATA_TRAIN, shuffle=True)
    (test_reviews, test_labels) = imdb.get_data(type=ImdbDataHandler.DATA_TEST, shuffle=True)

    # Simple bag of words with SGD
    sgd = SGDTextClassifier(train_reviews, train_labels,
                            test_texts=test_reviews, test_labels=test_labels,
                            compute_features=True)
    sgd.grid_search_cv(verbose=5, n_jobs=4)
    test_error = sgd.get_test_error()
    print "Test error in held out set: " + str(test_error)
    print "=" * 20

    # Now with bigrams too
    sgd = SGDTextClassifier(train_reviews, train_labels, ngram_range=(1,2),
                            test_texts=test_reviews, test_labels=test_labels,
                            compute_features=True)
    sgd.grid_search_cv(verbose=5, n_jobs=4)
    test_error = sgd.get_test_error()
    print "Test error in held out set: " + str(test_error)
    print "=" * 20

    # Simple bag of words with NB
    nb = NaiveBayesClassifier(train_reviews, train_labels,
                              test_texts=test_reviews, test_labels=test_labels)
    print "Loading data from original source"
    imdb = ImdbDataHandler(source=source)
    (train_reviews,
     train_labels) = imdb.get_data(type=ImdbDataHandler.DATA_TRAIN,
                                   shuffle=True)
    (test_reviews, test_labels) = imdb.get_data(type=ImdbDataHandler.DATA_TEST,
                                                shuffle=True)

    # Simple bag of words with SGD
    sgd = SGDTextClassifier(train_reviews,
                            train_labels,
                            test_texts=test_reviews,
                            test_labels=test_labels,
                            compute_features=True)
    sgd.grid_search_cv(verbose=5, n_jobs=4)
    test_error = sgd.get_test_error()
    print "Test error in held out set: " + str(test_error)
    print "=" * 20

    # Now with bigrams too
    sgd = SGDTextClassifier(train_reviews,
                            train_labels,
                            ngram_range=(1, 2),
                            test_texts=test_reviews,
                            test_labels=test_labels,
                            compute_features=True)
    sgd.grid_search_cv(verbose=5, n_jobs=4)
    test_error = sgd.get_test_error()
    print "Test error in held out set: " + str(test_error)
    print "=" * 20