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

Update files

Browse Source
This commit is contained in:
三洋三洋
2024-02-12 11:40:51 +00:00
parent fa6bcda721
commit 9bc165f955
4 changed files with 178 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -30,14 +30,14 @@ class TexTeller(VisionEncoderDecoderModel):
@classmethod
def from_pretrained(cls, model_path: str = None):
if model_path is None or model_path == cls.REPO_NAME:
if model_path is None or model_path == 'default':
return VisionEncoderDecoderModel.from_pretrained(cls.REPO_NAME)
model_path = Path(model_path).resolve()
return VisionEncoderDecoderModel.from_pretrained(str(model_path))
@classmethod
def get_tokenizer(cls, tokenizer_path: str = None) -> RobertaTokenizerFast:
if tokenizer_path is None or tokenizer_path == cls.REPO_NAME:
if tokenizer_path is None or tokenizer_path == 'default':
return RobertaTokenizerFast.from_pretrained(cls.REPO_NAME)
tokenizer_path = Path(tokenizer_path).resolve()
return RobertaTokenizerFast.from_pretrained(str(tokenizer_path))

8
src/start_web.sh
View File

@@ -1,9 +1,9 @@
#!/usr/bin/env bash
set -exu
export CHECKPOINT_DIR="OleehyO/TexTeller"
export TOKENIZER_DIR="OleehyO/TexTeller"
export USE_CUDA=True # True or False (case-sensitive)
export NUM_BEAM=5
export CHECKPOINT_DIR="default"
export TOKENIZER_DIR="default"
export USE_CUDA=False # True or False (case-sensitive)
export NUM_BEAM=1
streamlit run web.py