From 0d9051b7dd29b96ee2dd6e8c799594929df807c4 Mon Sep 17 00:00:00 2001 From: stephantul Date: Sat, 12 Sep 2026 12:03:51 +0200 Subject: [PATCH] docs: fix classifier docstring --- model2vec/train/classifier.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/model2vec/train/classifier.py b/model2vec/train/classifier.py index 779700e..adee335 100644 --- a/model2vec/train/classifier.py +++ b/model2vec/train/classifier.py @@ -173,8 +173,9 @@ def fit( :param device: The device to train on. If this is "auto", the device is chosen automatically. :param X_val: The texts to be used for validation. :param y_val: The labels to be used for validation. - :param class_weight: The weight of the classes. If None, all classes are weighted equally. Must - have the same length as the number of classes. + :param class_weight: The weight of the classes. If None, all classes are weighted equally. + If "balanced", weights are computed as the inverse class frequency. + If a dict, it must map each class to its weight. :param validation_steps: The number of steps to run validation for. If None, validation steps are estimated from the data. :param random_seed: The random seed to use. Defaults to 42. :return: The fitted model.