File tree Expand file tree Collapse file tree 1 file changed +6
-2
lines changed
Expand file tree Collapse file tree 1 file changed +6
-2
lines changed Original file line number Diff line number Diff line change 88from sentence_transformers import SentenceTransformerTrainer , losses
99from sentence_transformers .losses .BatchHardTripletLoss import BatchHardTripletLossDistanceFunction
1010from sentence_transformers .model_card import ModelCardCallback as STModelCardCallback
11- from sentence_transformers .training_args import BatchSamplers
11+ from sentence_transformers .training_args import BatchSamplers , SentenceTransformerTrainingArguments
1212from sklearn .preprocessing import LabelEncoder
1313from torch import nn
1414from transformers import __version__ as transformers_version
@@ -47,7 +47,11 @@ def __init__(
4747 self ._setfit_model = setfit_model
4848 self ._setfit_args = setfit_args
4949 self .logs_prefix = "embedding"
50- super ().__init__ (model = setfit_model .model_body , ** kwargs )
50+ super ().__init__ (
51+ model = setfit_model .model_body ,
52+ args = SentenceTransformerTrainingArguments (output_dir = setfit_args .output_dir ),
53+ ** kwargs ,
54+ )
5155 self ._apply_training_arguments (setfit_args )
5256
5357 for callback in list (self .callback_handler .callbacks ):
You can’t perform that action at this time.
0 commit comments