zuminghuang commited on
Commit
ad6d3bc
·
verified ·
1 Parent(s): 5d599d5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -37
app.py CHANGED
@@ -52,7 +52,7 @@ preset_prompts = [
52
 
53
 
54
  def send_pdf_to_parse(file_path, server_ip, port, route="/upload", api_key=None):
55
- url = f"http://{server_ip}:{port}{route}"
56
  headers = {}
57
  if api_key:
58
  headers["Authorization"] = f"Bearer {api_key}"
@@ -65,12 +65,12 @@ def send_pdf_to_parse(file_path, server_ip, port, route="/upload", api_key=None)
65
 
66
 
67
 
68
- async def send_pdf_async_aiohttp(file_path, server_ip, port, route="/upload", api_key=None):
69
  """使用aiohttp异步发送PDF"""
70
  url = f"http://{server_ip}:{port}{route}"
71
  headers = {}
72
- if api_key:
73
- headers["Authorization"] = f"Bearer {api_key}"
74
 
75
  try:
76
  async with aiohttp.ClientSession() as session:
@@ -100,38 +100,41 @@ IP = os.environ.get("IP")
100
 
101
  PORT = os.environ.get("PORT")
102
 
103
-
104
 
105
  client = AsyncOpenAI(
106
  api_key=openai_api_key,
107
- base_url=openai_api_base,
108
  )
109
 
110
 
111
  async def request(messages):
112
-
113
  chat_completion_from_base64 = await client.chat.completions.create(
114
  messages=messages,
 
 
 
115
  model="Qwen2_5VL",
116
  max_completion_tokens=4096,
117
  stream=True,
118
  temperature=0.0,
119
  top_p=0.95
120
  )
121
-
122
  page = ""
123
  async for chunk in chat_completion_from_base64:
124
  if chunk.choices[0].delta.content:
125
  content = chunk.choices[0].delta.content
126
-
127
  choice = chunk.choices[0]
128
  if choice.finish_reason is not None:
129
  print(f"end reason = {choice.finish_reason}")
130
  break
131
  page += content
132
-
133
  yield content
134
-
135
 
136
  def images_to_pdf(img_paths, pdf_path):
137
 
@@ -165,7 +168,7 @@ def encode_image(image_path):
165
  return base64.b64encode(image_file.read()).decode("utf-8")
166
 
167
  def build_message(image_path, prompt):
168
-
169
  content = [
170
  {
171
  "type": "image_url",
@@ -175,14 +178,14 @@ def build_message(image_path, prompt):
175
  },
176
  {"type": "text", 'text': prompt}
177
  ]
178
-
179
-
180
  messages = [
181
  {"role": "system", "content": "You are a helpful assistant."},
182
  {'role': 'user', 'content': content}
183
-
184
  ]
185
-
186
  return messages
187
 
188
 
@@ -211,14 +214,14 @@ async def doc_parser(doc_path, prompt):
211
  for idx, page in enumerate(pages, start=1):
212
  img_path = tmpdir / f"page_{idx}.png"
213
  page.save(img_path, "PNG")
214
-
215
  messages = build_message(img_path, prompt)
216
  queries.append(messages)
217
-
218
  else:
219
  messages = build_message(doc_path, prompt)
220
  queries.append(messages)
221
-
222
  all_pages = []
223
  all_pages_raw = []
224
  for query in queries:
@@ -231,19 +234,19 @@ async def doc_parser(doc_path, prompt):
231
  print(all_pages)
232
  yield "\n---\n".join(all_pages), "\n\n".join(all_pages_raw)
233
 
234
-
235
  def compress_directory_to_zip(directory_path, output_zip_path):
236
  try:
237
  with zipfile.ZipFile(output_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
238
 
239
-
240
  for root, dirs, files in os.walk(directory_path):
241
  for file in files:
242
-
243
  file_path = os.path.join(root, file)
244
 
245
  arcname = os.path.relpath(file_path, directory_path)
246
-
247
  zipf.write(file_path, arcname)
248
  return 0
249
  except Exception as e:
@@ -263,7 +266,7 @@ def check_prompt(prompt):
263
  return prompt
264
 
265
  def to_file(image_path):
266
-
267
  if image_path.endswith("Academic_Papers.png"):
268
  image_path = image_path.replace("Academic_Papers.png", "Academic_Papers.pdf")
269
 
@@ -293,7 +296,7 @@ async def process_file(file_path):
293
  images_to_pdf(file_path, tmp_file_path)
294
  else:
295
  tmp_file_path = file_path
296
- asyncio.create_task(send_pdf_async_aiohttp(tmp_file_path, IP, PORT))
297
 
298
  return str(tmp_file_path)
299
 
@@ -302,7 +305,7 @@ if __name__ == '__main__':
302
  with gr.Blocks() as demo:
303
  with gr.Row():
304
  with gr.Column(variant='panel', scale=5):
305
-
306
  file = gr.File(label='Please upload a PDF or image', file_types=['.pdf', '.png', '.jpeg', '.jpg'], type="filepath")
307
  prompts = gr.Dropdown(
308
  choices=preset_prompts,
@@ -320,7 +323,7 @@ if __name__ == '__main__':
320
  pdf_show = PDF(label='Preview', interactive=False, visible=True, height=800)
321
 
322
 
323
-
324
  example_root = os.path.join(os.path.dirname(__file__), 'examples')
325
  images = [
326
  os.path.join(example_root, f)
@@ -334,9 +337,9 @@ if __name__ == '__main__':
334
  file_path = [
335
  os.path.join(example_root, f)
336
  for f in ["Financial_Reports.png", "Books.png", "Magazines.png", "Academic_Papers.png"]
337
-
338
  ]
339
-
340
  with gr.Row():
341
  for i, label in enumerate(["Financial Reports(IMG)", "Books(IMG)", "Magazines(IMG)", "Academic Papers(PDF)"]):
342
  with gr.Column(scale=1, min_width=120):
@@ -348,11 +351,11 @@ if __name__ == '__main__':
348
  show_download_button=False
349
  )
350
  gr.Button(label).click(fn=to_file, inputs=gr.State(file_path[i]), outputs=file)
351
-
352
-
353
  download_btn = gr.Button("⬇️ Generate download link", size="sm")
354
  output_file = gr.File(label='Parse result', interactive=False, elem_id="down-file-box",visible=False)
355
-
356
  gr.HTML("""
357
  <style>
358
  #down-file-box {
@@ -367,9 +370,9 @@ if __name__ == '__main__':
367
  line_breaks=True)
368
  with gr.Tab('Markdown text'):
369
  md_text = gr.TextArea(lines=45, show_copy_button=True)
370
-
371
-
372
-
373
  file.change(fn=process_file, inputs=file, outputs=pdf_show)
374
 
375
 
@@ -386,9 +389,9 @@ if __name__ == '__main__':
386
  inputs=[file, prompts],
387
  outputs=[md, md_text]
388
  )
389
-
390
  clear_bu.add([file, md, pdf_show, md_text])
391
-
392
  download_btn.click(
393
  fn=download_markdown_file,
394
  inputs=md_text,
 
52
 
53
 
54
  def send_pdf_to_parse(file_path, server_ip, port, route="/upload", api_key=None):
55
+ url = f"{openai_api_base}{route}"
56
  headers = {}
57
  if api_key:
58
  headers["Authorization"] = f"Bearer {api_key}"
 
65
 
66
 
67
 
68
+ async def send_pdf_async_aiohttp(file_path, server_ip, route="/upload", Authorization=None):
69
  """使用aiohttp异步发送PDF"""
70
  url = f"http://{server_ip}:{port}{route}"
71
  headers = {}
72
+ if Authorization:
73
+ headers["Authorization"] = f"Bearer {Authorization}"
74
 
75
  try:
76
  async with aiohttp.ClientSession() as session:
 
100
 
101
  PORT = os.environ.get("PORT")
102
 
103
+ Authorization = os.environ.get("Authorization")
104
 
105
  client = AsyncOpenAI(
106
  api_key=openai_api_key,
107
+ base_url=openai_api_base + "/v1",
108
  )
109
 
110
 
111
  async def request(messages):
112
+
113
  chat_completion_from_base64 = await client.chat.completions.create(
114
  messages=messages,
115
+ extra_headers={
116
+ "Authorization": f"Bearer {Authorization}"
117
+ },
118
  model="Qwen2_5VL",
119
  max_completion_tokens=4096,
120
  stream=True,
121
  temperature=0.0,
122
  top_p=0.95
123
  )
124
+
125
  page = ""
126
  async for chunk in chat_completion_from_base64:
127
  if chunk.choices[0].delta.content:
128
  content = chunk.choices[0].delta.content
129
+
130
  choice = chunk.choices[0]
131
  if choice.finish_reason is not None:
132
  print(f"end reason = {choice.finish_reason}")
133
  break
134
  page += content
135
+
136
  yield content
137
+
138
 
139
  def images_to_pdf(img_paths, pdf_path):
140
 
 
168
  return base64.b64encode(image_file.read()).decode("utf-8")
169
 
170
  def build_message(image_path, prompt):
171
+
172
  content = [
173
  {
174
  "type": "image_url",
 
178
  },
179
  {"type": "text", 'text': prompt}
180
  ]
181
+
182
+
183
  messages = [
184
  {"role": "system", "content": "You are a helpful assistant."},
185
  {'role': 'user', 'content': content}
186
+
187
  ]
188
+
189
  return messages
190
 
191
 
 
214
  for idx, page in enumerate(pages, start=1):
215
  img_path = tmpdir / f"page_{idx}.png"
216
  page.save(img_path, "PNG")
217
+
218
  messages = build_message(img_path, prompt)
219
  queries.append(messages)
220
+
221
  else:
222
  messages = build_message(doc_path, prompt)
223
  queries.append(messages)
224
+
225
  all_pages = []
226
  all_pages_raw = []
227
  for query in queries:
 
234
  print(all_pages)
235
  yield "\n---\n".join(all_pages), "\n\n".join(all_pages_raw)
236
 
237
+
238
  def compress_directory_to_zip(directory_path, output_zip_path):
239
  try:
240
  with zipfile.ZipFile(output_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
241
 
242
+
243
  for root, dirs, files in os.walk(directory_path):
244
  for file in files:
245
+
246
  file_path = os.path.join(root, file)
247
 
248
  arcname = os.path.relpath(file_path, directory_path)
249
+
250
  zipf.write(file_path, arcname)
251
  return 0
252
  except Exception as e:
 
266
  return prompt
267
 
268
  def to_file(image_path):
269
+
270
  if image_path.endswith("Academic_Papers.png"):
271
  image_path = image_path.replace("Academic_Papers.png", "Academic_Papers.pdf")
272
 
 
296
  images_to_pdf(file_path, tmp_file_path)
297
  else:
298
  tmp_file_path = file_path
299
+ asyncio.create_task(send_pdf_async_aiohttp(tmp_file_path, server_ip=openai_api_base, Authorization=Authorization))
300
 
301
  return str(tmp_file_path)
302
 
 
305
  with gr.Blocks() as demo:
306
  with gr.Row():
307
  with gr.Column(variant='panel', scale=5):
308
+
309
  file = gr.File(label='Please upload a PDF or image', file_types=['.pdf', '.png', '.jpeg', '.jpg'], type="filepath")
310
  prompts = gr.Dropdown(
311
  choices=preset_prompts,
 
323
  pdf_show = PDF(label='Preview', interactive=False, visible=True, height=800)
324
 
325
 
326
+
327
  example_root = os.path.join(os.path.dirname(__file__), 'examples')
328
  images = [
329
  os.path.join(example_root, f)
 
337
  file_path = [
338
  os.path.join(example_root, f)
339
  for f in ["Financial_Reports.png", "Books.png", "Magazines.png", "Academic_Papers.png"]
340
+
341
  ]
342
+
343
  with gr.Row():
344
  for i, label in enumerate(["Financial Reports(IMG)", "Books(IMG)", "Magazines(IMG)", "Academic Papers(PDF)"]):
345
  with gr.Column(scale=1, min_width=120):
 
351
  show_download_button=False
352
  )
353
  gr.Button(label).click(fn=to_file, inputs=gr.State(file_path[i]), outputs=file)
354
+
355
+
356
  download_btn = gr.Button("⬇️ Generate download link", size="sm")
357
  output_file = gr.File(label='Parse result', interactive=False, elem_id="down-file-box",visible=False)
358
+
359
  gr.HTML("""
360
  <style>
361
  #down-file-box {
 
370
  line_breaks=True)
371
  with gr.Tab('Markdown text'):
372
  md_text = gr.TextArea(lines=45, show_copy_button=True)
373
+
374
+
375
+
376
  file.change(fn=process_file, inputs=file, outputs=pdf_show)
377
 
378
 
 
389
  inputs=[file, prompts],
390
  outputs=[md, md_text]
391
  )
392
+
393
  clear_bu.add([file, md, pdf_show, md_text])
394
+
395
  download_btn.click(
396
  fn=download_markdown_file,
397
  inputs=md_text,