TexTellerv2 release

src/models/ocr_model/model/TexTeller.py

@@ -17,7 +17,7 @@ from transformers import (
 
 
 class TexTeller(VisionEncoderDecoderModel):
-    REPO_NAME = '/home/lhy/code/TexTeller/src/models/ocr_model/train/train_result/TexTellerv2/checkpoint-588000'
+    REPO_NAME = 'OleehyO/TexTeller'
     def __init__(self, decoder_path=None, tokenizer_path=None):
         encoder = ViTModel(ViTConfig(
             image_size=FIXED_IMG_SIZE,

src/start_web.sh

@@ -3,8 +3,7 @@ set -exu
 
 export CHECKPOINT_DIR="default"
 export TOKENIZER_DIR="default"
-# export USE_CUDA=False  # True or False (case-sensitive)
-export USE_CUDA=True  # True or False (case-sensitive)
-export NUM_BEAM=10
+export USE_CUDA=False  # True or False (case-sensitive)
+export NUM_BEAM=1
 
 streamlit run web.py

src/web.py

@@ -65,7 +65,6 @@ tex = r'''
 def get_model():
     return TexTeller.from_pretrained(os.environ['CHECKPOINT_DIR'])
 
-
 @st.cache_resource
 def get_tokenizer():
     return TexTeller.get_tokenizer(os.environ['TOKENIZER_DIR'])