<jupyter_start><jupyter_text>Load data  IMDB<jupyter_code>def get_binary_imdb_data(base_dir, num_files=5000):
    def get_df(corpus_type):
        texts = []
        y = []
        corpus_type_dir = os.path.join(base_dir, f'{corpus_type}/pos/')
        for i, f_name in enumerate(os.listdir(corpus_type_dir)):
            if i == num_files: break
            f_path = os.path.join(corpus_type_dir, f_name)
            with open(f_path) as f:
                texts.append(f.read())
                y.append(1)
        corpus_type_dir = os.path.join(base_dir, f'{corpus_type}/neg/')
        for i, f_name in enumerate(os.listdir(corpus_type_dir)):
            if i == num_files: break
            f_path = os.path.join(corpus_type_dir, f_name)
            with open(f_path) as f:
                texts.append(f.read())
                y.append(0)
        df = pd.DataFrame()
        df['y'] = y
        df['text'] = texts
        return df
    train_data = shuffle(get_df('train'), random_state=0)
    test_data = shuffle(get_df('test'), random_state=0)
    return train_data, test_data
imdb_train, imdb_test = get_binary_imdb_data("data/aclImdb", num_files=None)<jupyter_output><empty_output><jupyter_text>Rotten tomatoes<jupyter_code>def get_binary_RT_data(base_dir):
    texts = []
    y = []
    f_path = os.path.join(base_dir, "rt-polarity.pos")
    with open(f_path, encoding = "ISO-8859-1") as f:
        for text in f:
            texts.append(text)
            y.append(1)
    f_path = os.path.join(base_dir, "rt-polarity.neg")
    with open(f_path, encoding = "ISO-8859-1") as f:
        for text in f:
            texts.append(text)
            y.append(0)
    X_train, X_test, y_train, y_test = train_test_split(
        texts, y, stratify=y, test_size=0.2, random_state=0
    )
    train_data = pd.DataFrame({"text": X_train, "y": y_train})
    test_data = pd.DataFrame({"text": X_test, "y": y_test})
    return train_data, test_data
rt_train, rt_test = get_binary_RT_data("data/rt-polaritydata")<jupyter_output><empty_output><jupyter_text>Airline sentiment<jupyter_code>def get_airline_binary_sentiment_data(path):
    df = pd.read_csv(path, encoding = "ISO-8859-1")
    df = df[df.airline_sentiment.isin(["negative","positive"])]
    df.loc[df.airline_sentiment == "negative","airline_sentiment"] = 0
    df.loc[df.airline_sentiment == "positive","airline_sentiment"] = 1
    X_train, X_test, y_train, y_test = train_test_split(
        df.text, df.airline_sentiment, stratify=df.airline_sentiment, test_size=0.2, random_state=0
    )
    train_data = pd.DataFrame({"text": X_train, "y": y_train})
    test_data = pd.DataFrame({"text": X_test, "y": y_test})
    return train_data, test_data
air_train, air_test = get_airline_binary_sentiment_data("data/Airline-Sentiment-2-w-AA.csv")<jupyter_output><empty_output><jupyter_text>Grid search<jupyter_code>transformers = [
    ("tf", TfidfTransformer(use_idf=False, sublinear_tf=True)),
    ("tfidf", TfidfTransformer(sublinear_tf=True)),
    ("tfpf", vectorizers.TfpfVectorizer(sublinear_tf=True)),
    ("tfrf", vectorizers.TfrfVectorizer(sublinear_tf=True)),
    ("tficf", vectorizers.TfIcfVectorizer(sublinear_tf=True)),
    ("tfbinicf", vectorizers.TfBinIcfVectorizer(sublinear_tf=True)),
    ("tfchi2", vectorizers.Tfchi2Vectorizer(sublinear_tf=True)),
    ("tfgr", vectorizers.TfgrVectorizer(sublinear_tf=True)),
    ("tfrrf", vectorizers.TfrrfVectorizer(sublinear_tf=True)),
    ("tfor", vectorizers.TforVectorizer(sublinear_tf=True))
]
param_grid = {
    "vectorizer__ngram_range": [(1, 1), (1, 2), (1, 3)],
    "vectorizer__max_features": [2000, 5000, None],
    "clf__C": [0.01, 0.1, 1, 10],
}
def grid_search(param_grid, train, test, n_jobs=-1):
    results = dict()
    
    for name, transformer in transformers:
        pipeline = Pipeline([
            ("vectorizer", CountVectorizer()),
            ("transformer", transformer),
            ("clf", LogisticRegression(solver="saga", random_state=0))
        ])
        search = GridSearchCV(pipeline, param_grid, cv=5, n_jobs=n_jobs, scoring="accuracy")
        search.fit(train.text, train.y)

        test_preds_proba = search.best_estimator_.predict_proba(test.text)
        test_preds = np.argmax(test_preds_proba, axis=1)

        train_accuracy = search.best_score_
        test_accuracy = accuracy_score(test.y, test_preds)
        f1 = f1_score(test.y, test_preds)
        print(f"{name:10} train accuracy: {train_accuracy:<15} test accuracy: {test_accuracy}")

        local_results = {
            "best_params": search.best_params_,
            "test_accuracy": test_accuracy,
            "f1": f1,
            "preds_proba": test_preds_proba[:, 1]
        }
        results[name] = local_results
    return results<jupyter_output><empty_output><jupyter_text>Correlation heatmap<jupyter_code>def plot_correlation_heatmap(results, title):
    probs = []
    labels = []
    for name, _ in transformers:
        probs.append(results[name]["preds_proba"])
        accuracy = results[name]["test_accuracy"]
        labels.append(f"{name}: {accuracy:.4f}")
    probs = np.array(probs)
    corr = np.corrcoef(probs)

    fig, ax = plt.subplots(figsize=(10,10))
    ax = sns.heatmap(corr, 
                     cmap="viridis", 
                     square=True, 
                     annot=corr, 
                     linewidth=0.5,
                     xticklabels=labels,
                     yticklabels=labels)
    ax.set_title(title)
    plt.show()
results = grid_search(param_grid, rt_train, rt_test, n_jobs=7)
plot_correlation_heatmap(results, "Rotten Tomatoes Binary Dataset")
results = grid_search(param_grid, imdb_train, imdb_test, n_jobs=4)
plot_correlation_heatmap(results, "IMDB Binary Sentiment Dataset")
results = grid_search(param_grid, air_train, air_test, n_jobs=7)
plot_correlation_heatmap(results, "Airline Sentiment Binary Dataset")<jupyter_output>tf         train accuracy: 0.9221187175043327 test accuracy: 0.9194456474664356
tfidf      train accuracy: 0.9243934142114385 test accuracy: 0.9259419662191425
tfpf       train accuracy: 0.9226603119584056 test accuracy: 0.9211779991338241
tfrf       train accuracy: 0.9219020797227037 test accuracy: 0.9194456474664356
tficf      train accuracy: 0.9236351819757366 test accuracy: 0.9198787353832828
tfbinicf   train accuracy: 0.9221187175043327 test accuracy: 0.9194456474664356
tfchi2     train accuracy: 0.9129116117850953 test accuracy: 0.9107838891294933
tfgr       train accuracy: 0.8996967071057193 test accuracy: 0.895625812039844
tfrrf      train accuracy: 0.9222270363951474 test accuracy: 0.9194456474664356
tfor       train accuracy: 0.9246100519930676 test accuracy: 0.920744911216977
{'fraction_non_alphanumeric': 0.12233009708737864, 'fraction_numerical': 0.06273338312173264, 'mean_word_length': 3.493959731543624, 'pattern_counts': {'":': 15, '<': 21, '<?xml version=': 0, '>': 20, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '568324', 'license': 'MIT', 'n_tokens_mistral': 2674, 'n_tokens_neox': 2390, 'n_words': 404, 'path': 'examples/binary_comparison.ipynb', 'repo_name': 'bernardoduarte/textvec', 'size': '6695'}