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

Update model config

Browse Source
This commit is contained in:
三洋三洋
2024-06-22 22:08:08 +08:00
parent c5e859517a
commit 36a2680d28
2 changed files with 27 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/models/ocr_model/model/TexTeller.py
View File

@@ -18,11 +18,11 @@ from transformers import (
class TexTeller(VisionEncoderDecoderModel):
REPO_NAME = 'OleehyO/TexTeller'
def __init__(self):
config = VisionEncoderDecoderConfig.from_pretrained('/home/lhy/code/TexTeller/src/models/ocr_model/model/trocr-small')
config.encoder.image_size = FIXED_IMG_SIZE
config.encoder.num_channels = IMG_CHANNELS
config.decoder.vocab_size=VOCAB_SIZE
config.decoder.max_position_embeddings=MAX_TOKEN_SIZE
config = VisionEncoderDecoderConfig.from_pretrained(Path(__file__).resolve().parent / "config.json")
config.encoder.image_size = FIXED_IMG_SIZE
config.encoder.num_channels = IMG_CHANNELS
config.decoder.vocab_size = VOCAB_SIZE
config.decoder.max_position_embeddings = MAX_TOKEN_SIZE
super().__init__(config=config)