linjieccc commited on
Commit
66a930c
1 Parent(s): 2cc2c6f

add docprompt_v2

Browse files
Files changed (1) hide show
  1. app.py +10 -7
app.py CHANGED
@@ -248,21 +248,19 @@ def get_base64(path):
248
  return base64_str
249
 
250
 
251
- def process_prompt(prompt, document, lang="ch"):
252
  if not prompt:
253
  prompt = "What is the total actual and/or obligated expenses of ECG Center?"
254
  if document is None:
255
  return None, None, None
256
 
257
  access_token = os.environ['token']
258
-
259
  url = f"https://aip.baidubce.com/rpc/2.0/nlp-itec/poc/docprompt?access_token={access_token}"
260
 
261
  base64_str = get_base64(document)
262
 
263
- r = requests.post(url, json={"doc": base64_str, "prompt": [prompt], "lang": lang})
264
  response = r.json()
265
-
266
  predictions = response['result']
267
  img_list = response['image']
268
  pages = [Image.open(BytesIO(base64.b64decode(img))) for img in img_list]
@@ -283,7 +281,7 @@ def load_example_document(img, prompt):
283
  if img is not None:
284
  document = prompt_files[prompt]
285
  lang = lang_map[document]
286
- preview, answer, answer_text = process_prompt(prompt, document, lang)
287
  return document, prompt, preview, gr.update(visible=True), answer, answer_text
288
  else:
289
  return None, None, None, gr.update(visible=False), None, None
@@ -441,6 +439,11 @@ with gr.Blocks(css=CSS) as demo:
441
  value="en",
442
  label="Select OCR Language (Please choose ch for Chinese images.)",
443
  )
 
 
 
 
 
444
 
445
  with gr.Row():
446
  clear_button = gr.Button("Clear", variant="secondary")
@@ -495,13 +498,13 @@ with gr.Blocks(css=CSS) as demo:
495
 
496
  prompt.submit(
497
  fn=process_prompt,
498
- inputs=[prompt, document, ocr_lang],
499
  outputs=[image, output, output_text],
500
  )
501
 
502
  submit_button.click(
503
  fn=process_prompt,
504
- inputs=[prompt, document, ocr_lang],
505
  outputs=[image, output, output_text],
506
  )
507
 
 
248
  return base64_str
249
 
250
 
251
+ def process_prompt(prompt, document, lang="ch", model="docprompt_v1"):
252
  if not prompt:
253
  prompt = "What is the total actual and/or obligated expenses of ECG Center?"
254
  if document is None:
255
  return None, None, None
256
 
257
  access_token = os.environ['token']
 
258
  url = f"https://aip.baidubce.com/rpc/2.0/nlp-itec/poc/docprompt?access_token={access_token}"
259
 
260
  base64_str = get_base64(document)
261
 
262
+ r = requests.post(url, json={"doc": base64_str, "prompt": [prompt], "lang": lang, "model": model})
263
  response = r.json()
 
264
  predictions = response['result']
265
  img_list = response['image']
266
  pages = [Image.open(BytesIO(base64.b64decode(img))) for img in img_list]
 
281
  if img is not None:
282
  document = prompt_files[prompt]
283
  lang = lang_map[document]
284
+ preview, answer, answer_text = process_prompt(prompt, document, lang, "docprompt_v1")
285
  return document, prompt, preview, gr.update(visible=True), answer, answer_text
286
  else:
287
  return None, None, None, gr.update(visible=False), None, None
 
439
  value="en",
440
  label="Select OCR Language (Please choose ch for Chinese images.)",
441
  )
442
+ model = gr.Radio(
443
+ choices=["docprompt_v1", "docprompt_v2"],
444
+ value="docprompt_v1",
445
+ label="Select Inference Model.",
446
+ )
447
 
448
  with gr.Row():
449
  clear_button = gr.Button("Clear", variant="secondary")
 
498
 
499
  prompt.submit(
500
  fn=process_prompt,
501
+ inputs=[prompt, document, ocr_lang, model],
502
  outputs=[image, output, output_text],
503
  )
504
 
505
  submit_button.click(
506
  fn=process_prompt,
507
+ inputs=[prompt, document, ocr_lang, model],
508
  outputs=[image, output, output_text],
509
  )
510