YogeLiu/TexTeller
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

TexTellerv2

Browse Source
This commit is contained in:
三洋三洋
2024-03-25 11:46:43 +00:00
parent a42df1510f
commit ef7cccff03
4 changed files with 23 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/models/ocr_model/train/train.py
View File

@@ -38,7 +38,7 @@ def train(model, tokenizer, train_dataset, eval_dataset, collate_fn_with_tokeniz
)
# trainer.train(resume_from_checkpoint=None)
trainer.train(resume_from_checkpoint='/home/lhy/code/TexTeller/src/models/ocr_model/train/train_result/TexTellerv2/checkpoint-64000')
trainer.train(resume_from_checkpoint='/home/lhy/code/TexTeller/src/models/ocr_model/train/train_result/TexTellerv2/checkpoint-288000')
def evaluate(model, tokenizer, eval_dataset, collate_fn):
@@ -94,9 +94,9 @@ if __name__ == '__main__':
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/TexTeller/src/models/ocr_model/train/train_result/TexTellerv2/checkpoint-64000')
model = TexTeller.from_pretrained('/home/lhy/code/TexTeller/src/models/ocr_model/train/train_result/TexTellerv2/checkpoint-588000')
enable_train = True
enable_train = False
enable_evaluate = True
if enable_train:
train(model, tokenizer, train_dataset, eval_dataset, collate_fn_with_tokenizer)