Merge pull request #53 from SeanLee97/bugfix/trainer

Bugfix/trainer

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
 files = setup.py angle_emb/__init__.py
-current_version = 0.3.4
+current_version = 0.3.5
 commit = True
 tag = True

angle_emb/__init__.py

@@ -3,4 +3,4 @@
 from .angle import *
 
 
-__version__ = '0.3.4'
+__version__ = '0.3.5'

angle_emb/angle.py

@@ -744,6 +744,7 @@ def __init__(self,
  pooling_strategy='all',
  padding_strategy=self.pooler.padding_strategy,
  is_llm=False)
+ self.kl_loss_fct = nn.KLDivLoss(reduction='batchmean')
  logger.info(f'Train with alignment, teacher={fixed_teacher_name_or_path}')
 
  def compute_loss(self, model, inputs, return_outputs=False):

angle_emb/train_cli.py

@@ -166,14 +166,17 @@ def main():
  argument_kwargs['report_to'] = 'wandb'
 
  trainer_kwargs = None
- if args.apply_tdmse:
+ if args.fixed_teacher_name_or_path is not None:
  trainer_kwargs = {
+ 'fixed_teacher_name_or_path': args.fixed_teacher_name_or_path
+ }
+ if args.apply_tdmse:
+ trainer_kwargs = dict(trainer_kwargs, **{
  'apply_tdmse_kl': args.apply_tdmse_kl,
  'tdmse_kl_temperature': args.tdmse_kl_temperature,
  'tdmse_teacher_lambda': args.tdmse_teacher_lambda,
  'tdmse_student_lambda': args.tdmse_student_lambda,
- 'fixed_teacher_name_or_path': args.fixed_teacher_name_or_path,
- }
+ })
 
  model.fit(
  train_ds=train_ds,

setup.py

@@ -16,7 +16,7 @@
 
 setup(
  name='angle_emb',
- version='0.3.4',
+ version='0.3.5',
  description='AnglE-optimize Text Embeddings',
  long_description=readme,
  long_description_content_type="text/markdown",