Upload folder using huggingface_hub

filt_result_by_bleu.py

@@ -0,0 +1,9 @@
+from datasets import load_dataset
+import pandas as pd
+
+decoding_res = '/home/zychen/hwproject/my_modeling_phase_1/mytest_3600_test5k/decoding_res.json'
+dataset2 = load_dataset("json", data_files=decoding_res)["train"]
+print(f"Number of examples: {len(dataset2)}")
+decoding_df = dataset2.to_pandas()
+
+decoding_df

make_comet_hyp_and_ref.py

@@ -0,0 +1,38 @@
+from datasets import load_dataset
+import pandas as pd
+
+text_src_jsonl = '/home/zychen/hwproject/my_modeling_phase_1/mytest/text_src.jsonl'
+dataset = load_dataset("json", data_files=text_src_jsonl)["train"]
+print(f"Number of examples: {len(dataset)}")
+text_src_df = dataset.to_pandas()
+
+decoding_res = '/home/zychen/hwproject/my_modeling_phase_1/mytest_3600_test5k/decoding_res.json'
+dataset2 = load_dataset("json", data_files=decoding_res)["train"]
+print(f"Number of examples: {len(dataset2)}")
+decoding_df = dataset2.to_pandas()
+
+df_merged = pd.concat([text_src_df, decoding_df], axis=1)
+print(df_merged.columns.tolist(), df_merged.iloc[4500])
+
+
+def clean(sentence):
+    return ''.join(sentence.split())
+
+
+df = df_merged
+with open('text_src.txt', 'w', encoding='utf-8') as f:
+    for text in df['text_src']:
+        # cleaned_text = clean(text)
+        f.write(text + '\n')
+
+# 将trans_res_seg列的内容写入hyp.txt
+with open('hyp.txt', 'w', encoding='utf-8') as f:
+    for text in df['trans_res_seg']:
+        cleaned_text = clean(text)
+        f.write(cleaned_text + '\n')
+
+# 将gt_seg列的内容写入ref.txt
+with open('ref.txt', 'w', encoding='utf-8') as f:
+    for text in df['gt_seg']:
+        cleaned_text = clean(text)
+        f.write(cleaned_text + '\n')

make_jsonl.py

@@ -0,0 +1,98 @@
+import json
+import csv
+import os
+from tqdm import tqdm  # 导入tqdm用于显示进度条
+import re
+from utils import *
+import traceback
+
+
+def process_json_files(csv_path, output_dir):
+    # 创建输出目录
+    os.makedirs(output_dir, exist_ok=True)
+    json_file = open(os.path.join(output_dir, 'output1.jsonl'),
+                     'w',
+                     encoding='utf-8')
+    try:
+        # 读取CSV文件
+        with open(csv_path, 'r', encoding='utf-8') as csv_file:
+            csv_reader = csv.reader(csv_file)
+            next(csv_reader)  # 跳过标题行
+
+            # 使用tqdm包装csv_reader以显示进度条
+            for row in tqdm(csv_reader,
+                            desc="Processing JSON files",
+                            unit="file"):
+                json_path = row[0]  # 获取JSON文件路径
+                # print('row', row)
+                try:
+                    # 读取JSON文件
+                    with open(json_path, 'r', encoding='utf-8') as f:
+                        json_data = json.load(f)
+                    img_path = row[1]
+                    shape = cv2.imread(img_path).shape
+                    #element -> tuple: (word_text, word_bbox, normed_word_bbox)
+                    #  resize_box(text_['src_word_bboxes'][i],shape)
+                    doc_triplet = []
+                    doc_tgt_sen_trans = []
+                    doc_words_boxes_list = []
+                    # 处理JSON数据
+                    for key, value in json_data.items():
+                        if value.get("attribute") == 'text_block':
+                            for text_ in value.get('text', []):
+                                combined_list = [(
+                                    text_['src_words'][i],
+                                    text_['src_word_bboxes'][i],
+                                ) for i in range(len(text_['src_words']))]
+                                doc_words_boxes_list.extend(combined_list)
+                                # print(f'combined_list:{combined_list}')
+                                doc_tgt_sen_trans.append(
+                                    text_['tgt_text.zh-CN'])
+                    processed_list = [
+                        (src_w, src_w_boxes, resize_box(src_w_boxes, shape))
+                        for (src_w, src_w_boxes) in doc_words_boxes_list
+                    ]
+                    # print(f'processed:{processed_list}')
+                    sorted_tuple_list = tblr_reading_order_detector(
+                        processed_list)
+
+                    text_src_list = [atuple[0] for atuple in sorted_tuple_list]
+                    layout_src_list = [
+                        atuple[2] for atuple in sorted_tuple_list
+                    ]
+                    text_src = ' '.join(text_src_list)
+                    tgt_sen_trans = ''.join(doc_tgt_sen_trans)
+                    # print('text_src', text_src)
+                    data_dict = {
+                        "img_path": img_path,
+                        "text_src": text_src,
+                        "layout_src": layout_src_list,
+                        "tgt_sen_trans": tgt_sen_trans
+                    }
+                    # print(data_dict)
+                    json_line = json.dumps(data_dict, ensure_ascii=False)
+                    json_file.write(json_line + '\n')
+
+                except FileNotFoundError:
+                    print(f"File not found: {json_path}")
+                except json.JSONDecodeError:
+                    print(f"Error decoding JSON in file: {json_path}")
+                except KeyError as e:
+                    print(f"Missing key {e} in file: {json_path}")
+                except Exception as e:
+                    print(f"Unexpected error processing {json_path}: {str(e)}")
+                    traceback.print_exc()
+
+    except FileNotFoundError:
+        print(f"CSV file not found: {csv_path}")
+    except Exception as e:
+        print(f"Error reading CSV file: {str(e)}")
+
+    print("Processing completed!")
+
+
+# csv_path = '/home/zychen/hwproject/my_modeling_phase_1/dataset/output_part2.csv'  # 替换为你的CSV文件路径
+csv_path = '/home/zychen/hwproject/my_modeling_phase_1/dataset/output.csv'  # 替换为你的CSV文件路径
+output_dir = '/home/zychen/hwproject/my_modeling_phase_1/dataset'  # 输出目录名
+
+process_json_files(csv_path, output_dir)

make_text_src_list.py

@@ -0,0 +1,43 @@
+import json
+
+from model_and_train import MyDataset, prepare_dataset_df, prepare_tokenizer
+from torch.utils.data import DataLoader, Dataset
+
+dataset_dir = "/home/zychen/hwproject/my_modeling_phase_1/dataset"
+data_file = f"{dataset_dir}/testset_10k.jsonl"
+if __name__ == "__main__":
+
+    encoder_ckpt_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/lilt-roberta-en-base"
+
+    tgt_tokenizer_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/bert-base-chinese-tokenizer"
+
+    src_tokenizer, tgt_tokenizer = prepare_tokenizer(
+        src_tokenizer_dir=encoder_ckpt_dir,
+        tgt_tokenizer_dir=tgt_tokenizer_dir,
+    )
+    dataset_df = prepare_dataset_df(data_file=data_file)
+    my_dataset = MyDataset(df=dataset_df,
+                           src_tokenizer=src_tokenizer,
+                           tgt_tokenizer=tgt_tokenizer,
+                           max_src_length=512,
+                           max_target_length=512)
+    print(len(my_dataset))
+    from torch.utils.data import Subset
+    num_test = 5000  #total 10k
+    my_dataset = Subset(my_dataset, range(0, num_test))
+    # my_dataloader = DataLoader(
+    #     my_dataset,
+    #     batch_size=batch_size,
+    #     shuffle=False,
+    # )
+    img_name_list = dataset_df["img_path"].iloc[0:num_test].tolist()
+    text_src_list = dataset_df["text_src"].iloc[0:num_test].tolist()
+    with open('./mytest/text_src.jsonl', "w") as decoding_res_file:
+        for img_name, text_src in zip(img_name_list, text_src_list):
+            res_dict = {
+                "img_name": img_name,
+                "text_src": text_src,
+            }
+
+            record = f"{json.dumps(res_dict, ensure_ascii=False)}\n"
+            decoding_res_file.write(record)

model_and_train.py

@@ -0,0 +1,293 @@
+# basic imports
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "4"
+
+# other external imports
+import pandas as pd
+# torch imports
+import torch
+from datasets import load_dataset
+from torch.utils.data import DataLoader, Dataset
+# transformers imports
+from transformers import (BertConfig, BertTokenizer, EncoderDecoderConfig,
+                          EncoderDecoderModel, LayoutLMv3Tokenizer, LiltConfig,
+                          LiltModel, Seq2SeqTrainer, Seq2SeqTrainingArguments,
+                          default_data_collator)
+
+# internal imports
+
+
+
+# prepare tokenizer.
+def prepare_tokenizer(src_tokenizer_dir, tgt_tokenizer_dir):
+    src_tokenizer = LayoutLMv3Tokenizer.from_pretrained(src_tokenizer_dir)
+    tgt_tokenizer = BertTokenizer.from_pretrained(tgt_tokenizer_dir)
+
+    return src_tokenizer, tgt_tokenizer
+
+
+# read data points.
+def prepare_dataset_df(data_file):
+
+    def filter_fn(exam):
+        bboxes = exam["layout_src"]
+        for box in bboxes:
+            x0, y0, x1, y1 = box
+            if (x0 > x1) or (y0 > y1):
+                print("(x0 > x1) or (y0 > y1)")
+                return False
+            for cor in box:
+                if cor < 0 or cor > 1000:
+                    # print("cor < 0 or cor > 1000")
+                    # print(exam['img_path'],box)
+                    return False
+        return True
+
+    dataset = load_dataset("json", data_files=data_file)["train"]
+    print()
+    print(f"Number of examples: {len(dataset)}")
+    print()
+
+    dataset = dataset.filter(filter_fn, num_proc=48)
+
+    dataset_df = dataset.to_pandas()
+    # dataset_df = pd.read_json(data_file, lines=True, orient="records")
+
+    # filter the nan data points.
+    dataset_df = dataset_df[~dataset_df["tgt_sen_trans"].isna()]
+    dataset_df = dataset_df[~dataset_df["text_src"].isna()]
+    dataset_df = dataset_df[~dataset_df["layout_src"].isna()]
+    # remove entries where "text_src" length is less than 3
+    dataset_df = dataset_df[dataset_df["text_src"].str.len() >= 3]
+    # reconstruct the idx to avoid index_error.
+    dataset_df = dataset_df.reset_index(drop=True)
+
+    print(f"Number of examples after filtered: {len(dataset_df)}")
+    return dataset_df
+
+
+class MyDataset(Dataset):
+
+    def __init__(
+        self,
+        df,
+        src_tokenizer,
+        tgt_tokenizer,
+        max_src_length,
+        max_target_length,
+    ):
+        self.df = df
+        self.src_tokenizer = src_tokenizer
+        self.tgt_tokenizer = tgt_tokenizer
+        self.max_src_length = max_src_length
+        self.max_target_length = max_target_length
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        # get text_src + layout_src + tgt_trans.
+        text_src = self.df['text_src'][idx]
+        layout_src = self.df['layout_src'][idx]
+        tgt_trans = self.df['tgt_sen_trans'][idx]
+
+        # read in annotations at word-level (words, word boxes)
+        words_ = text_src.split(" ")
+        word_boxes_ = layout_src
+        # print('words', words_, len(words_), len(word_boxes_))
+        assert len(words_) == len(word_boxes_)
+        words = []
+        word_boxes = []
+        for word, word_box in zip(words_, word_boxes_):
+            if (word_box[0] >= word_box[2]) or (word_box[1] >= word_box[3]):
+                continue
+
+            words.append(word)
+            word_boxes.append(word_box)
+
+        assert len(words) == len(word_boxes)
+
+        encoding = self.src_tokenizer(
+            words,
+            boxes=word_boxes,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_src_length,
+        )
+
+        # construct labels.
+        labels = self.tgt_tokenizer(
+            tgt_trans,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_target_length)["input_ids"]
+        # important: make sure that PAD tokens are ignored by the loss function
+        labels = [
+            label if label != self.tgt_tokenizer.pad_token_id else -100
+            for label in labels
+        ]
+
+        encoding["labels"] = labels
+
+        assert len(encoding['input_ids']) == self.max_src_length
+        assert len(encoding['attention_mask']) == self.max_src_length
+        assert len(encoding['bbox']) == self.max_src_length
+        assert len(encoding['labels']) == self.max_target_length
+
+        # finally, convert everything to PyTorch tensors
+        for k, v in encoding.items():
+            encoding[k] = torch.as_tensor(encoding[k])
+
+        return encoding
+
+
+def prepare_model(src_tokenizer,
+                  tgt_tokenizer,
+                  max_src_len,
+                  max_tgt_len,
+                  num_encoder_hidden_layers,
+                  num_decoder_hidden_layers,
+                  encoder_ckpt_dir,
+                  model_ckpt_dir=None):
+    config_encoder = LiltConfig.from_pretrained(
+        encoder_ckpt_dir,
+        max_position_embeddings=max_src_len + 2,
+        num_hidden_layers=num_encoder_hidden_layers)
+    config_decoder = BertConfig(vocab_size=tgt_tokenizer.vocab_size,
+                                max_position_embeddings=max_tgt_len,
+                                num_hidden_layers=num_decoder_hidden_layers)
+
+    model_config = EncoderDecoderConfig.from_encoder_decoder_configs(
+        encoder_config=config_encoder,
+        decoder_config=config_decoder,
+    )
+    model = EncoderDecoderModel(config=model_config, )
+
+    model.config.decoder_start_token_id = tgt_tokenizer.cls_token_id
+    model.config.pad_token_id = tgt_tokenizer.pad_token_id
+    model.config.vocab_size = tgt_tokenizer.vocab_size
+    model.config.eos_token_id = tgt_tokenizer.pad_token_id
+
+    from safetensors.torch import load_file
+    if model_ckpt_dir:
+        bin_path = f"{model_ckpt_dir}/pytorch_model.bin"
+        safetensors_path = f"{model_ckpt_dir}/model.safetensors"
+        if os.path.exists(bin_path):
+            state_dict = torch.load(bin_path)
+        elif os.path.exists(safetensors_path):
+            state_dict = load_file(safetensors_path)
+        else:
+            raise FileNotFoundError(
+                "Neither pytorch_model.bin nor model.safetensors found in the specified directory."
+            )
+        model.load_state_dict(state_dict, strict=False)
+        model.save_pretrained(
+            f"continued_{model_ckpt_dir}")  #save at continued training
+    else:
+        # Loading the pre-trained params and then save the model, including its configuration.
+        tmp_encoder = LiltModel.from_pretrained(
+            pretrained_model_name_or_path=encoder_ckpt_dir,
+            config=config_encoder,
+        )
+        # tmp_encoder = LiltModel(config=config_encoder)
+        model.encoder = tmp_encoder
+        # model.save_pretrained("undertrained_default_safe_true")
+        model.save_pretrained("undertrained_safe_serialization_False", safe_serialization=False)
+        # model.load_state_dict(torch.load(f"undertrained/pytorch_model.bin"))
+
+        bin_path = "undertrained_safe_serialization_False/pytorch_model.bin"
+        safetensors_path = "undertrained_default_safe_true/model.safetensors"
+        if os.path.exists(bin_path):
+            state_dict = torch.load(bin_path)
+        elif os.path.exists(safetensors_path):
+            state_dict = load_file(safetensors_path)
+        else:
+            raise FileNotFoundError(
+                "Neither pytorch_model.bin nor model.safetensors found in the specified directory."
+            )
+        model.load_state_dict(state_dict, strict=False)
+
+    print(model.config)
+    print(model)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    # hyper-parameters.
+    ## for model.
+    MAX_TGT_LEN = 512
+    MAX_SRC_LEN = 512
+    num_encoder_hidden_layers = 12
+    num_decoder_hidden_layers = 12
+
+    ## for training.
+    num_instances = 500000  #total 620082 ./dataset/merged.jsonl Number of examples after filtered: 547084
+    learning_rate = 1e-4
+    batch_size = 28
+    num_train_steps = 400000 #400000
+    output_dir = f"./train.lr_{learning_rate}.bsz_{batch_size}.step_{num_train_steps}.layer_{num_encoder_hidden_layers}-{num_decoder_hidden_layers}"
+    save_total_limit = 100
+    save_steps = num_train_steps // save_total_limit
+
+    dataset_dir = "/home/zychen/hwproject/my_modeling_phase_1/dataset"
+    data_file = f"{dataset_dir}/merged.jsonl"
+
+    # model_ckpt_dir = '/home/zychen/hwproject/my_modeling_phase_1/train.lr_0.0001.bsz_8.step_400000.layer_12-12/checkpoint-32000'
+    model_ckpt_dir = '/home/zychen/hwproject/my_modeling_phase_1/train.lr_0.0001.bsz_16.step_500000.layer_12-12_36k+20k/checkpoint-20000'
+    encoder_ckpt_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/lilt-roberta-en-base"
+
+    tgt_tokenizer_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/bert-base-chinese-tokenizer"
+
+    src_tokenizer, tgt_tokenizer = prepare_tokenizer(
+        src_tokenizer_dir=encoder_ckpt_dir,
+        tgt_tokenizer_dir=tgt_tokenizer_dir,
+    )
+    dataset_df = prepare_dataset_df(data_file=data_file)[:num_instances]
+    print(f"\nnum_instances: {len(dataset_df)}\n")
+    print(dataset_df)
+    my_dataset = MyDataset(
+        df=dataset_df,
+        src_tokenizer=src_tokenizer,
+        tgt_tokenizer=tgt_tokenizer,
+        max_src_length=MAX_SRC_LEN,
+        max_target_length=MAX_TGT_LEN,
+    )
+    model = prepare_model(src_tokenizer=src_tokenizer,
+                          tgt_tokenizer=tgt_tokenizer,
+                          max_src_len=MAX_SRC_LEN,
+                          max_tgt_len=MAX_TGT_LEN,
+                          num_encoder_hidden_layers=num_encoder_hidden_layers,
+                          num_decoder_hidden_layers=num_decoder_hidden_layers,
+                          encoder_ckpt_dir=encoder_ckpt_dir,
+                          model_ckpt_dir=model_ckpt_dir)
+
+    training_args = Seq2SeqTrainingArguments(
+        predict_with_generate=False,
+        evaluation_strategy="no",
+        per_device_train_batch_size=batch_size,
+        fp16=True,
+        output_dir=output_dir,
+        logging_steps=1,
+        # save_strategy="epoch",
+        learning_rate=learning_rate,
+        max_steps=num_train_steps,
+        warmup_ratio=0.05,
+        save_total_limit=save_total_limit,
+        save_steps=save_steps,
+        save_safetensors=False,
+    )
+    # print(training_args)
+    # instantiate trainer
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        compute_metrics=None,
+        train_dataset=my_dataset,
+        eval_dataset=None,
+        data_collator=default_data_collator,
+    )
+
+    trainer.train()

old_model_and_train.py

@@ -0,0 +1,266 @@
+# basic imports
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+# transformers imports
+from transformers import LiltConfig, BertConfig, EncoderDecoderConfig, EncoderDecoderModel, BertTokenizer, LayoutLMv3Tokenizer, LiltModel
+from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
+from transformers import default_data_collator
+from datasets import load_dataset
+
+# torch imports
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+# internal imports
+
+# other external imports
+import pandas as pd
+
+
+# prepare tokenizer.
+def prepare_tokenizer(src_tokenizer_dir, tgt_tokenizer_dir):
+    src_tokenizer = LayoutLMv3Tokenizer.from_pretrained(src_tokenizer_dir)
+    tgt_tokenizer = BertTokenizer.from_pretrained(tgt_tokenizer_dir)
+
+    return src_tokenizer, tgt_tokenizer
+
+
+# read data points.
+def prepare_dataset_df(data_file):
+
+    def filter_fn(exam):
+        bboxes = exam['block_list']
+        for box in bboxes:
+            x0, y0, x1, y1 = box["block_bbox"]
+            if (x0 > x1) or (y0 > y1):
+                print(box["block_bbox"])
+                return False
+            for cor in box["block_bbox"]:
+                # if cor < 0 or cor > 1000:
+                if cor <0:
+                    return False
+        return True
+
+    dataset = load_dataset("json", data_files=data_file)["train"]
+    print()
+    print(f"Number of examples: {len(dataset)}")
+    print()
+    # print(dataset[0]['block_list'])
+    dataset = dataset.filter(filter_fn, num_proc=48)
+
+    dataset_df = dataset.to_pandas()
+    # dataset_df = pd.read_json(data_file, lines=True, orient="records")
+
+    # filter the nan data points.
+    # dataset_df = dataset_df[~dataset_df["tgt_sen_trans"].isna()]
+    # dataset_df = dataset_df[~dataset_df["text_src"].isna()]
+    # dataset_df = dataset_df[~dataset_df["layout_src"].isna()]
+
+    # reconstruct the idx to avoid index_error.
+    dataset_df = dataset_df.reset_index(drop=True)
+
+    print(f"Number of examples after filtered: {len(dataset_df)}")
+    print(dataset_df)
+    return dataset_df
+
+
+class MyDataset(Dataset):
+
+    def __init__(
+        self,
+        df,
+        src_tokenizer,
+        tgt_tokenizer,
+        max_src_length,
+        max_target_length,
+    ):
+        self.df = df
+        self.src_tokenizer = src_tokenizer
+        self.tgt_tokenizer = tgt_tokenizer
+        self.max_src_length = max_src_length
+        self.max_target_length = max_target_length
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        # get text_src + layout_src + tgt_trans.
+        text_src = self.df['text_src'][idx]
+        layout_src = self.df['layout_src'][idx]
+        tgt_trans = self.df['tgt_sen_trans'][idx]
+
+        # read in annotations at word-level (words, word boxes)
+        words_ = text_src.split(" ")
+        word_boxes_ = layout_src
+        assert len(words_) == len(word_boxes_)
+        words = []
+        word_boxes = []
+        for word, word_box in zip(words_, word_boxes_):
+            if (word_box[0] >= word_box[2]) or (word_box[1] >= word_box[3]):
+                continue
+
+            words.append(word)
+            word_boxes.append(word_box)
+
+        assert len(words) == len(word_boxes)
+
+        encoding = self.src_tokenizer(
+            words,
+            boxes=word_boxes,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_src_length,
+        )
+
+        # construct labels.
+        labels = self.tgt_tokenizer(
+            tgt_trans,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_target_length)["input_ids"]
+        # important: make sure that PAD tokens are ignored by the loss function
+        labels = [
+            label if label != self.tgt_tokenizer.pad_token_id else -100
+            for label in labels
+        ]
+
+        encoding["labels"] = labels
+
+        assert len(encoding['input_ids']) == self.max_src_length
+        assert len(encoding['attention_mask']) == self.max_src_length
+        assert len(encoding['bbox']) == self.max_src_length
+        assert len(encoding['labels']) == self.max_target_length
+
+        # finally, convert everything to PyTorch tensors
+        for k, v in encoding.items():
+            encoding[k] = torch.as_tensor(encoding[k])
+
+        return encoding
+
+
+def prepare_model(src_tokenizer,
+                  tgt_tokenizer,
+                  max_src_len,
+                  max_tgt_len,
+                  num_encoder_hidden_layers,
+                  num_decoder_hidden_layers,
+                  encoder_ckpt_dir,
+                  model_ckpt_dir=None):
+    config_encoder = LiltConfig.from_pretrained(
+        encoder_ckpt_dir,
+        max_position_embeddings=max_src_len + 2,
+        num_hidden_layers=num_encoder_hidden_layers)
+    config_decoder = BertConfig(vocab_size=tgt_tokenizer.vocab_size,
+                                max_position_embeddings=max_tgt_len,
+                                num_hidden_layers=num_decoder_hidden_layers)
+
+    model_config = EncoderDecoderConfig.from_encoder_decoder_configs(
+        encoder_config=config_encoder,
+        decoder_config=config_decoder,
+    )
+    model = EncoderDecoderModel(config=model_config, )
+
+    model.config.decoder_start_token_id = tgt_tokenizer.cls_token_id
+    model.config.pad_token_id = tgt_tokenizer.pad_token_id
+    model.config.vocab_size = tgt_tokenizer.vocab_size
+    model.config.eos_token_id = tgt_tokenizer.pad_token_id
+
+    if model_ckpt_dir:
+        model.load_state_dict(
+            torch.load(f"{model_ckpt_dir}/pytorch_model.bin"))
+    else:
+        # Loading the pre-trained params and then save the model, including its configuration.
+        tmp_encoder = LiltModel.from_pretrained(
+            pretrained_model_name_or_path=encoder_ckpt_dir,
+            config=config_encoder,
+        )
+        # tmp_encoder = LiltModel(config=config_encoder)
+        model.encoder = tmp_encoder
+        model.save_pretrained("undertrained")
+        model.load_state_dict(torch.load(f"undertrained/pytorch_model.bin"))
+
+    print(model.config)
+    print(model)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    # hyper-parameters.
+    ## for model.
+    MAX_TGT_LEN = 512
+    MAX_SRC_LEN = 512
+    num_encoder_hidden_layers = 12
+    num_decoder_hidden_layers = 12
+
+    ## for training.
+    # wc 12420 ./dataset/scene_imgs/jsons/en_json/en_scene.jsonl
+    # wc 12230 ./dataset/scene_imgs/jsons/zh_json/zh_scene.jsonl
+    num_instances = 500000
+    learning_rate = 1e-4
+    batch_size = 16
+    num_train_steps = 400000
+    output_dir = f"./train.lr_{learning_rate}.bsz_{batch_size}.step_{num_train_steps}.layer_{num_encoder_hidden_layers}-{num_decoder_hidden_layers}"
+    save_total_limit = 100
+    save_steps = num_train_steps // save_total_limit
+
+    # dataset_dir = "/home/zychen/hwproject/my_modeling_phase_1/dataset/scene_imgs/jsons/en_json/en_scene.jsonl"
+    data_file = "/home/zychen/hwproject/my_modeling_phase_1/dataset/scene_imgs/jsons/en_json/en_scene.jsonl"
+
+    model_ckpt_dir = None
+
+    encoder_ckpt_dir = "./Tokenizer_PretrainedWeights/lilt-roberta-en-base"
+
+    tgt_tokenizer_dir = "./Tokenizer_PretrainedWeights/bert-base-chinese-tokenizer"
+
+    src_tokenizer, tgt_tokenizer = prepare_tokenizer(
+        src_tokenizer_dir=encoder_ckpt_dir,
+        tgt_tokenizer_dir=tgt_tokenizer_dir,
+    )
+    dataset_df = prepare_dataset_df(data_file=data_file)[:num_instances]
+    print(f"\nnum_instances: {len(dataset_df)}\n")
+    my_dataset = MyDataset(
+        df=dataset_df,
+        src_tokenizer=src_tokenizer,
+        tgt_tokenizer=tgt_tokenizer,
+        max_src_length=MAX_SRC_LEN,
+        max_target_length=MAX_TGT_LEN,
+    )
+    model = prepare_model(src_tokenizer=src_tokenizer,
+                          tgt_tokenizer=tgt_tokenizer,
+                          max_src_len=MAX_SRC_LEN,
+                          max_tgt_len=MAX_TGT_LEN,
+                          num_encoder_hidden_layers=num_encoder_hidden_layers,
+                          num_decoder_hidden_layers=num_decoder_hidden_layers,
+                          encoder_ckpt_dir=encoder_ckpt_dir,
+                          model_ckpt_dir=model_ckpt_dir)
+
+    training_args = Seq2SeqTrainingArguments(
+        predict_with_generate=False,
+        evaluation_strategy="no",
+        per_device_train_batch_size=batch_size,
+        fp16=True,
+        output_dir=output_dir,
+        logging_steps=1,
+        # save_strategy="epoch",
+        learning_rate=learning_rate,
+        max_steps=num_train_steps,
+        warmup_ratio=0.05,
+        save_total_limit=save_total_limit,
+        save_steps=save_steps,
+    )
+
+    # instantiate trainer
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        compute_metrics=None,
+        train_dataset=my_dataset,
+        eval_dataset=None,
+        data_collator=default_data_collator,
+    )
+
+    trainer.train()

sample_generate.py

@@ -0,0 +1,82 @@
+# basic imports
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+# other external imports
+import pandas as pd
+# torch imports
+import torch
+from datasets import load_dataset
+from torch.utils.data import DataLoader, Dataset
+# transformers imports
+from transformers import (BertConfig, BertTokenizer, EncoderDecoderConfig,
+                          EncoderDecoderModel, LayoutLMv3Tokenizer, LiltConfig,
+                          LiltModel, Seq2SeqTrainer, Seq2SeqTrainingArguments,
+                          default_data_collator)
+
+# internal imports
+
+
+
+def prepare_tokenizer(src_tokenizer_dir, tgt_tokenizer_dir):
+    src_tokenizer = LayoutLMv3Tokenizer.from_pretrained(src_tokenizer_dir)
+    tgt_tokenizer = BertTokenizer.from_pretrained(tgt_tokenizer_dir)
+
+    return src_tokenizer, tgt_tokenizer
+
+
+if __name__ == "__main__":
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = 'cpu'
+    print(device)
+    checkpoints_dir = '/home/zychen/hwproject/my_modeling_phase_1/train.lr_0.0001.bsz_8.step_400000.layer_12-12_36000'
+    model = EncoderDecoderModel.from_pretrained(
+        f"{checkpoints_dir}/checkpoint-36000").to(device)
+    encoder_ckpt_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/lilt-roberta-en-base"
+    tgt_tokenizer_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/bert-base-chinese-tokenizer"
+
+    src_tokenizer, tgt_tokenizer = prepare_tokenizer(
+        src_tokenizer_dir=encoder_ckpt_dir,
+        tgt_tokenizer_dir=tgt_tokenizer_dir,
+    )
+    model.eval()
+
+    from model_and_train import (MyDataset, prepare_dataset_df,
+                                 prepare_tokenizer)
+
+    dataset_dir = "/home/zychen/hwproject/my_modeling_phase_1/dataset"
+    data_file = f"{dataset_dir}/merged.jsonl"
+    dataset_df = prepare_dataset_df(data_file=data_file)[:1000]
+    print(f"\nnum_instances: {len(dataset_df)}\n")
+    print(dataset_df)
+    my_dataset = MyDataset(
+        df=dataset_df,
+        src_tokenizer=src_tokenizer,
+        tgt_tokenizer=tgt_tokenizer,
+        max_src_length=512,
+        max_target_length=512,
+    )
+    sample = my_dataset[0]
+    from transformers import GenerationConfig
+    generation_config = GenerationConfig(
+        max_length=512,
+        early_stopping=True,
+        num_beams=1,
+        use_cache=True,
+        length_penalty=1.0,
+    )
+
+    with torch.no_grad():
+        generation_config = None
+        outputs = model.generate(
+            input_ids=sample['input_ids'].unsqueeze(
+                0),  # 添加 unsqueeze 以增加 batch 维度
+            attention_mask=sample['attention_mask'].unsqueeze(0),
+            do_sample=False,
+            generation_config=generation_config,
+            bos_token_id=0)
+        decoded_preds = tgt_tokenizer.batch_decode(outputs,
+                                                   skip_special_tokens=True)
+        print(decoded_preds)
+        print(sample['labels'])

test_bleu.py

@@ -0,0 +1,100 @@
+# basic imports
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+# other external imports
+import pandas as pd
+import sacrebleu
+# torch imports
+import torch
+from datasets import load_dataset
+from torch.utils.data import DataLoader, Dataset
+# transformers imports
+from tqdm import tqdm
+from transformers import (BertConfig, BertTokenizer, EncoderDecoderConfig,
+                          EncoderDecoderModel, LayoutLMv3Tokenizer, LiltConfig,
+                          LiltModel, Seq2SeqTrainer, Seq2SeqTrainingArguments,
+                          default_data_collator)
+
+# internal imports
+
+
+
+def prepare_tokenizer(src_tokenizer_dir, tgt_tokenizer_dir):
+    src_tokenizer = LayoutLMv3Tokenizer.from_pretrained(src_tokenizer_dir)
+    tgt_tokenizer = BertTokenizer.from_pretrained(tgt_tokenizer_dir)
+
+    return src_tokenizer, tgt_tokenizer
+
+
+def prepare_dataset_df(data_file):
+    dataset_df = pd.read_json(data_file, lines=True)
+    return dataset_df
+
+
+if __name__ == "__main__":
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = 'cpu'
+    print(device)
+    checkpoints_dir = '/home/zychen/hwproject/my_modeling_phase_1/train.lr_0.0001.bsz_28.step_400000.layer_12-12'
+    model = EncoderDecoderModel.from_pretrained(
+        f"{checkpoints_dir}/checkpoint-64000").to(device)
+    encoder_ckpt_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/lilt-roberta-en-base"
+    tgt_tokenizer_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/bert-base-chinese-tokenizer"
+
+    src_tokenizer, tgt_tokenizer = prepare_tokenizer(
+        src_tokenizer_dir=encoder_ckpt_dir,
+        tgt_tokenizer_dir=tgt_tokenizer_dir,
+    )
+    model.eval()
+
+    dataset_dir = "/home/zychen/hwproject/my_modeling_phase_1/dataset"
+    data_file = f"{dataset_dir}/merged.jsonl"
+    dataset_df = prepare_dataset_df(data_file=data_file)[:5000]
+    print(f"\nnum_instances: {len(dataset_df)}\n")
+    from model_and_train import (MyDataset, prepare_dataset_df,
+                                 prepare_tokenizer)
+
+    my_dataset = MyDataset(
+        df=dataset_df,
+        src_tokenizer=src_tokenizer,
+        tgt_tokenizer=tgt_tokenizer,
+        max_src_length=512,
+        max_target_length=512,
+    )
+
+    dataloader = DataLoader(my_dataset, batch_size=4, shuffle=False)
+
+    references = []
+    predictions = []
+
+    with torch.no_grad():
+        for batch in tqdm(dataloader):
+            input_ids = batch['input_ids'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+            labels = batch['labels'].tolist()
+            outputs = model.generate(input_ids=input_ids,
+                                     attention_mask=attention_mask,
+                                     do_sample=True,
+                                     max_length=512,
+                                     num_beams=1,
+                                     use_cache=True,
+                                     length_penalty=1.0,
+                                     bos_token_id=0)
+
+            decoded_preds = tgt_tokenizer.batch_decode(
+                outputs, skip_special_tokens=True)
+            decoded_labels = tgt_tokenizer.batch_decode(
+                labels, skip_special_tokens=True)
+
+            predictions.extend(decoded_preds)
+            references.extend([label.split(' ') for label in decoded_labels])
+
+            predictions_str = ''.join(predictions)
+            references_str = ''.join([''.join(ref) for ref in references])
+
+            print(predictions_str, references_str)
+
+    bleu_score = sacrebleu.corpus_bleu(predictions, [references])
+    print(f"BLEU score: {bleu_score.score}")

test_bleu_chrf.py

@@ -0,0 +1,164 @@
+# basic imports
+import json
+import os
+
+import jieba
+# other external imports
+import pandas as pd
+# torch imports
+import torch
+# internal imports
+from model_and_train import MyDataset, prepare_dataset_df, prepare_tokenizer
+from nltk.translate.bleu_score import sentence_bleu
+from sacrebleu.metrics import CHRF
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+# transformers imports
+from transformers import BertTokenizer, EncoderDecoderModel
+
+chrf = CHRF(word_order=2)  # word_order=2 to be chrf++.
+
+os.environ["CUDA_VISIBLE_DEVICES"] = "5"
+
+# hyper-parameters.
+## for model.
+MAX_TGT_LEN = 512
+MAX_SRC_LEN = 512
+
+## for decoding.
+output_dir = "./mytest"
+os.makedirs(output_dir, exist_ok=True)
+early_stopping = True
+num_beams = 2
+length_penalty = 1.0
+batch_size = 16
+metric_res_filepath = os.path.join(output_dir, "metric_res.json")
+decoding_res_filepath = os.path.join(output_dir, "decoding_res.json")
+trained_model_dir = "/home/zychen/hwproject/my_modeling_phase_1/train.lr_0.0001.bsz_28.step_400000.layer_12-12/checkpoint-64000"
+
+dataset_dir = "/home/zychen/hwproject/my_modeling_phase_1/dataset"
+data_file = f"{dataset_dir}/testset_10k.jsonl"
+
+
+def no_blank(sen):
+    return "".join(sen.split())
+
+
+if __name__ == "__main__":
+
+    encoder_ckpt_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/lilt-roberta-en-base"
+
+    tgt_tokenizer_dir = "/home/zychen/hwproject/my_modeling_phase_1/Tokenizer_PretrainedWeights/bert-base-chinese-tokenizer"
+
+    src_tokenizer, tgt_tokenizer = prepare_tokenizer(
+        src_tokenizer_dir=encoder_ckpt_dir,
+        tgt_tokenizer_dir=tgt_tokenizer_dir,
+    )
+    dataset_df = prepare_dataset_df(data_file=data_file)
+    my_dataset = MyDataset(df=dataset_df,
+                           src_tokenizer=src_tokenizer,
+                           tgt_tokenizer=tgt_tokenizer,
+                           max_src_length=512,
+                           max_target_length=512)
+    print(len(my_dataset))
+    from torch.utils.data import Subset
+    num_test = 5000  #total 10k
+    my_dataset = Subset(my_dataset, range(0, num_test))
+    my_dataloader = DataLoader(
+        my_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+    )
+
+    # loading model and config from pretrained folder
+    model = EncoderDecoderModel.from_pretrained(trained_model_dir)
+    # device='cpu'
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    model.eval()
+
+    print(model)
+
+    # decoding testset
+    pred_res_list = []
+    gt_list = []
+
+    for batch in tqdm(my_dataloader):
+        # predict use generate
+        with torch.no_grad():
+            encoder_outputs = model.encoder(
+                input_ids=batch["input_ids"].to(device),
+                bbox=batch["bbox"].to(device),
+                attention_mask=batch["attention_mask"].to(device),
+            )
+            outputs = model.generate(
+                input_ids=batch["input_ids"].to(device),
+                attention_mask=batch["attention_mask"].to(device),
+                encoder_outputs=encoder_outputs,
+                max_length=MAX_TGT_LEN,
+                early_stopping=early_stopping,
+                num_beams=num_beams,
+                length_penalty=length_penalty,
+                use_cache=True,
+                decoder_start_token_id=0)
+
+        # decode
+        pred_str = tgt_tokenizer.batch_decode(outputs,
+                                              skip_special_tokens=True)
+        labels = batch["labels"]
+        labels[labels == -100] = tgt_tokenizer.pad_token_id
+        label_str = tgt_tokenizer.batch_decode(labels,
+                                               skip_special_tokens=True)
+
+        pred_res_list += pred_str
+        gt_list += label_str
+
+    gt_list = [no_blank(sen) for sen in gt_list]
+    pred_res_list = [no_blank(sen) for sen in pred_res_list]
+
+    # write the decoding res and compute metric.
+    img_name_list = dataset_df["img_path"].iloc[0:num_test].tolist()
+    text_src_list = dataset_df["text_src"].iloc[0:num_test].tolist()
+    bleu_list = []
+    chrf_list = []
+
+    pred_res_seg_list = [" ".join(jieba.cut(item)) for item in pred_res_list]
+    gt_seg_list = [" ".join(jieba.cut(item)) for item in gt_list]
+    print(len(text_src_list), len(pred_res_seg_list), len(gt_seg_list))
+    # print(img_name_list, pred_res_list, gt_seg_list)
+    assert len(img_name_list) == len(pred_res_seg_list) == len(gt_seg_list)
+
+    with open(decoding_res_filepath, "w") as decoding_res_file:
+        for img_name, text_src, pred_res_seg, gt_seg in zip(
+                img_name_list, text_src_list, pred_res_seg_list, gt_seg_list):
+
+            instance_bleu = sentence_bleu([gt_seg.split()],
+                                          pred_res_seg.split())
+            bleu_list.append(instance_bleu)
+
+            instance_chrf = chrf.sentence_score(
+                hypothesis=pred_res_seg,
+                references=[gt_seg],
+            ).score
+            chrf_list.append(instance_chrf)
+
+            res_dict = {
+                "img_name": img_name,
+                "text_src": text_src,
+                "instance_bleu": instance_bleu,
+                "instance_chrf": instance_chrf,
+                "trans_res_seg": pred_res_seg,
+                "gt_seg": gt_seg,
+            }
+
+            record = f"{json.dumps(res_dict, ensure_ascii=False)}\n"
+            decoding_res_file.write(record)
+
+    trans_avg_bleu = sum(bleu_list) / len(bleu_list)
+    trans_avg_chrf = sum(chrf_list) / len(chrf_list)
+    with open(metric_res_filepath, "w") as metric_res_file:
+        eval_res_dict = {
+            "trans_avg_bleu": trans_avg_bleu,
+            "trans_avg_chrf": trans_avg_chrf,
+        }
+        json.dump(eval_res_dict, metric_res_file, indent=4, ensure_ascii=False)

utils.py

@@ -0,0 +1,56 @@
+import numpy as np
+import cv2
+import os
+import json
+from functools import cmp_to_key
+
+def resize_box(box, ori_img_shape, nomarlized_img_shape=(1000, 1000, 3)):
+    """
+        box: [x0, y0, x1, y1],
+        ori_img_shape: 形如(1560, 1103, 3)
+    """
+
+    height_ratio = nomarlized_img_shape[0] / ori_img_shape[0]
+    width_ratio = nomarlized_img_shape[1] / ori_img_shape[1]
+
+    x0, y0, x1, y1 = box
+    norm_x0, norm_x1 = round(x0 * width_ratio), round(x1 * width_ratio)
+    norm_y0, norm_y1 = round(y0 * height_ratio), round(y1 * height_ratio)
+
+    return [norm_x0, norm_y0, norm_x1, norm_y1]
+
+def tblr_reading_order_detector(tuple_list):
+    """rule: top-to-bottom, left-to-right
+
+        tuple: (word_text, word_bbox, normed_word_bbox)
+        
+        return: sorted_tuple_list
+    """
+
+    def sort_cmp_fn(word_box1, word_box2):
+        """
+            sorted function的排序的2个元素的比较准则。
+            1. 比较box1和box2的y坐标，如果二者的高重合度达到了二者的50%，则位于同一行，否则位于不同行。
+            2. 如果位于同一行，那么比较二者的x0，如果box1_x0 < box2_x0，则返回-1，表示box_1<box_2，否则返回0（表示相等）或者1（box1>box2）。
+            3. 如果不位于同一行，那么比较二者的y0，如果box1_y0 < box2_y0，则返回-1，否则返回0或者1.
+        """
+
+        x0, y0, x1, y1 = word_box1[1]
+        x0_, y0_, x1_, y1_ = word_box2[1]
+
+        
+        if y0 < y0_:
+            return -1
+        elif y0 > y0_:
+            return 1
+        elif y0 == y0_:
+            if x0 <= x0_:
+                return -1
+            elif x0 > x0_:
+                return 1
+
+    sorted_tuple_list = sorted(tuple_list, key=cmp_to_key(sort_cmp_fn))
+    # print(sorted_word_box_list)
+
+    return sorted_tuple_list
+