修复了bug：当样本中出现非常长的公式(对应的token数可能超过2048)，会导致给label进行embedding时index out of range

src/models/ocr_model/train/train.py

@@ -9,13 +9,24 @@ from transformers import Trainer, TrainingArguments, Seq2SeqTrainer, Seq2SeqTrai
 
 from .training_args import CONFIG
 from ..model.TexTeller import TexTeller
-from ..utils.functional import tokenize_fn, collate_fn, img_transform_fn
+from ..utils.functional import tokenize_fn, collate_fn, img_transform_fn, filter_fn
 from ..utils.metrics import bleu_metric
 from ...globals import MAX_TOKEN_SIZE, MIN_WIDTH, MIN_HEIGHT   
 
 
 def train(model, tokenizer, train_dataset, eval_dataset, collate_fn_with_tokenizer):
     training_args = TrainingArguments(**CONFIG)
+    debug_mode = True
+    if debug_mode:
+        training_args.auto_find_batch_size = False
+        training_args.num_train_epochs = 2
+        training_args.per_device_train_batch_size = 2
+        training_args.per_device_eval_batch_size = 2 * training_args.per_device_train_batch_size
+        training_args.jit_mode_eval = False
+        training_args.torch_compile = False
+        training_args.dataloader_num_workers = 1
+
+    
     trainer = Trainer(
         model,
         training_args,
@@ -34,7 +45,7 @@ def evaluate(model, tokenizer, eval_dataset, collate_fn):
     eval_config = CONFIG.copy()
     eval_config['predict_with_generate'] = True
     generate_config = GenerationConfig(
-        max_new_tokens=MAX_TOKEN_SIZE,
+        max_length=MAX_TOKEN_SIZE-200,
         num_beams=1,
         do_sample=False,
         pad_token_id=tokenizer.pad_token_id,
@@ -67,23 +78,25 @@ if __name__ == '__main__':
     dataset = load_dataset(
         '/home/lhy/code/TexTeller/src/models/ocr_model/train/data/loader.py'
     )['train']
-    dataset = dataset.filter(lambda x: x['image'].height > MIN_HEIGHT and x['image'].width > MIN_WIDTH)
+    tokenizer = TexTeller.get_tokenizer('/home/lhy/code/TexTeller/src/models/tokenizer/roberta-tokenizer-7Mformulas')
+    filter_fn_with_tokenizer = partial(filter_fn, tokenizer=tokenizer)
+
+    # dataset = dataset.filter(lambda x: x['image'].height > MIN_HEIGHT and x['image'].width > MIN_WIDTH)
+    dataset = dataset.filter(filter_fn_with_tokenizer, num_proc=16)
     dataset = dataset.shuffle(seed=42)
     dataset = dataset.flatten_indices()
 
-    tokenizer = TexTeller.get_tokenizer('/home/lhy/code/TeXify/src/models/tokenizer/roberta-tokenizer-550Kformulas')
-
     map_fn = partial(tokenize_fn, tokenizer=tokenizer)
     tokenized_dataset = dataset.map(map_fn, batched=True, remove_columns=dataset.column_names, num_proc=8, load_from_cache_file=True)
     tokenized_dataset = tokenized_dataset.with_transform(img_transform_fn)
 
-    split_dataset = tokenized_dataset.train_test_split(test_size=0.05, seed=42)    
+    split_dataset = tokenized_dataset.train_test_split(test_size=0.005, seed=42)    
     train_dataset, eval_dataset = split_dataset['train'], split_dataset['test']
     collate_fn_with_tokenizer = partial(collate_fn, tokenizer=tokenizer)
     model = TexTeller()
     # model = TexTeller.from_pretrained('/home/lhy/code/TeXify/src/models/ocr_model/train/train_result/train_with_random_resize/checkpoint-80000')
 
-    enable_train = False
+    enable_train = True
     enable_evaluate = True
     if enable_train:
         train(model, tokenizer, train_dataset, eval_dataset, collate_fn_with_tokenizer)