add docprompt_v2

app.py

@@ -248,21 +248,19 @@ def get_base64(path):
     return base64_str
 
 
-def process_prompt(prompt, document, lang="ch"):
+def process_prompt(prompt, document, lang="ch", model="docprompt_v1"):
     if not prompt:
         prompt = "What is the total actual and/or obligated expenses of ECG Center?"
     if document is None:
         return None, None, None
 
     access_token = os.environ['token']
-
     url = f"https://aip.baidubce.com/rpc/2.0/nlp-itec/poc/docprompt?access_token={access_token}"
     
     base64_str = get_base64(document)
 
-    r = requests.post(url, json={"doc": base64_str, "prompt": [prompt], "lang": lang})
+    r = requests.post(url, json={"doc": base64_str, "prompt": [prompt], "lang": lang, "model": model})
     response = r.json()
-    
     predictions = response['result']
     img_list = response['image']
     pages = [Image.open(BytesIO(base64.b64decode(img))) for img in img_list]
@@ -283,7 +281,7 @@ def load_example_document(img, prompt):
     if img is not None:
         document = prompt_files[prompt]
         lang = lang_map[document]
-        preview, answer, answer_text = process_prompt(prompt, document, lang)
+        preview, answer, answer_text = process_prompt(prompt, document, lang, "docprompt_v1")
         return document, prompt, preview, gr.update(visible=True), answer, answer_text
     else:
         return None, None, None, gr.update(visible=False), None, None
@@ -441,6 +439,11 @@ with gr.Blocks(css=CSS) as demo:
                 value="en",
                 label="Select OCR Language (Please choose ch for Chinese images.)",
             )
+            model = gr.Radio(
+                choices=["docprompt_v1", "docprompt_v2"],
+                value="docprompt_v1",
+                label="Select Inference Model.",
+            )
 
             with gr.Row():
                 clear_button = gr.Button("Clear", variant="secondary")
@@ -495,13 +498,13 @@ with gr.Blocks(css=CSS) as demo:
 
     prompt.submit(
         fn=process_prompt,
-        inputs=[prompt, document, ocr_lang],
+        inputs=[prompt, document, ocr_lang, model],
         outputs=[image, output, output_text],
     )
 
     submit_button.click(
         fn=process_prompt,
-        inputs=[prompt, document, ocr_lang],
+        inputs=[prompt, document, ocr_lang, model],
         outputs=[image, output, output_text],
     )