fasttext_text_classification_train_save_predict_evaluate.py

1import fasttext
2
3# 1. Train a supervised classification model
4# Note: 'data.train.txt' should contain one example per line in the format:
5# __label__category_name some text content
6model = fasttext.train_supervised(input="cooking.train")
7
8# 2. Save the trained model
9model.save_model("model_cooking.bin")
10
11# 3. Predict the label of a new sentence
12# The first value in the result is the label, the second is the probability
13prediction = model.predict("Which organic food should I buy?")
14print(prediction)
15
16# 4. Evaluate the model on a test set
17# Returns a tuple: (number of examples, precision at 1, recall at 1)
18result = model.test("cooking.valid")
19print(f"Number of examples: {result[0]}")
20print(f"P@1: {result[1]}")
21print(f"R@1: {result[2]}")