add selector for cache usage

app.py

@@ -11,7 +11,7 @@ def process_filename(filename, question):
     return process_image(image)
 
 
-def process_image(image, question):
+def process_image(set_use_cache, image, question):
     repo_id = "naver-clova-ix/donut-base-finetuned-docvqa"
     print(f"Model repo: {repo_id}")
     processor = DonutProcessor.from_pretrained(repo_id)
@@ -33,7 +33,7 @@ def process_image(image, question):
         max_length=model.decoder.config.max_position_embeddings,
         pad_token_id=processor.tokenizer.pad_token_id,
         eos_token_id=processor.tokenizer.eos_token_id,
-        use_cache=False,
+        use_cache=set_use_cache,
         bad_words_ids=[[processor.tokenizer.unk_token_id]],
         return_dict_in_generate=True,
     )
@@ -48,17 +48,17 @@ def process_image(image, question):
     sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
     print(processor.token2json(sequence))
 
-    return [True, processor.token2json(sequence)['answer'], ""]
-
-def process_document(image, question):
-    ret = process_image(image, question)
-    return ret[1]
+    return processor.token2json(sequence)['answer']
 
 description = "DocVQA (document visual question answering)"
 
 demo = gr.Interface(
-    fn=process_document,
-    inputs=["image", gr.Textbox(label = "Question" )],
+    fn=process_image,
+    inputs=[
+        gr.Radio(["True", "False"], label="Use cache", info="Define if model.generate() should use cache"), 
+        "image", 
+        gr.Textbox(label = "Question" )
+    ],
     outputs=gr.Textbox(label = "Response" ),
     title="Extract data from image",
     description=description,