Updated Chipper model

app.py

@@ -13,7 +13,7 @@ from logits_ngrams import NoRepeatNGramLogitsProcessor, get_table_token_ids
 def run_prediction(sample, model, processor, mode):
 
     skip_tokens = get_table_token_ids(processor)
-    no_repeat_ngram_size = 10
+    no_repeat_ngram_size = 15
 
     if mode == "OCR":
         prompt = "<s><s_pretraining>"
@@ -35,9 +35,9 @@ def run_prediction(sample, model, processor, mode):
             decoder_input_ids=processor.tokenizer(prompt, add_special_tokens=False, return_tensors="pt").input_ids.to(device),
             logits_processor=[NoRepeatNGramLogitsProcessor(no_repeat_ngram_size, skip_tokens)],
             do_sample=True,
-            top_p=0.92, #.92,
+            top_p=0.92, 
             top_k=5,
-            no_repeat_ngram_size=0,
+            no_repeat_ngram_size=25,
             num_beams=3,
             output_attentions=False,
             output_hidden_states=False,
@@ -81,7 +81,7 @@ else:
 st.image(image, caption='Your target document')
 
 with st.spinner(f'Processing the document ...'):
-        pre_trained_model = "unstructuredio/chipper-fast-fine-tuning"
+        pre_trained_model = "unstructuredio/chipper-fast-fine-tuning-oct-23-release"
         processor = DonutProcessor.from_pretrained(pre_trained_model, token=os.environ['HF_TOKEN'])
         
         device = "cuda" if torch.cuda.is_available() else "cpu"

logits_ngrams.py

@@ -59,5 +59,5 @@ def _calc_banned_tokens(prev_input_ids, num_hypos, no_repeat_ngram_size, cur_len
 
 
 def get_table_token_ids(processor):
-    skip_tokens = {token_id for token, token_id in processor.tokenizer.get_added_vocab().items() if token.startswith("<t") or token.startswith("</t") }
+    return {token_id for token, token_id in processor.tokenizer.get_added_vocab().items() if token.startswith("<t") or token.startswith("</t") }