# Source: content/notes/nlp/text-classification.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.model_selection import GroupShuffleSplit
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.pipeline import make_pipeline
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import f1_score
texts, labels, groups = [], [], []
for group in range(30):
    texts.extend([f"billing payment invoice account{group}", f"technical crash error account{group}"])
    labels.extend([0, 1])
    groups.extend([group, group])
texts, labels, groups = np.array(texts), np.array(labels), np.array(groups)
train, test = next(GroupShuffleSplit(n_splits=1, test_size=.3, random_state=3).split(texts, labels, groups))
assert not set(groups[train]) & set(groups[test])
model = make_pipeline(TfidfVectorizer(ngram_range=(1, 2)), LogisticRegression(random_state=0))
model.fit(texts[train], labels[train])
prediction = model.predict(texts[test])
assert f1_score(labels[test], prediction, average="macro") == 1
vocab = model.named_steps["tfidfvectorizer"].vocabulary_
assert f"account{groups[test][0]}" not in vocab
print("grouped split, train-only vocabulary and baseline predictions passed")