openfree commited on
Commit
16a2fde
ยท
verified ยท
1 Parent(s): 00039aa

Update app-backup.py

Browse files
Files changed (1) hide show
  1. app-backup.py +275 -257
app-backup.py CHANGED
@@ -8,9 +8,11 @@ import re
8
  import uuid
9
  import pymupdf
10
 
 
 
 
11
  os.system('pip uninstall -y magic-pdf')
12
  os.system('pip install git+https://github.com/opendatalab/MinerU.git@dev')
13
-
14
  os.system('wget https://github.com/opendatalab/MinerU/raw/dev/scripts/download_models_hf.py -O download_models_hf.py')
15
  os.system('python download_models_hf.py')
16
 
@@ -26,7 +28,6 @@ with open('/home/user/magic-pdf.json', 'w') as file:
26
  json.dump(data, file, indent=4)
27
 
28
  os.system('cp -r paddleocr /home/user/.paddleocr')
29
- from gradio_pdf import PDF
30
 
31
  import gradio as gr
32
  from loguru import logger
@@ -35,23 +36,33 @@ from magic_pdf.data.data_reader_writer import FileBasedDataReader
35
  from magic_pdf.libs.hash_utils import compute_sha256
36
  from magic_pdf.tools.common import do_parse, prepare_env
37
 
 
 
 
 
38
  def create_css():
 
 
 
39
  return """
40
- /* ์ „์ฒด ์Šคํƒ€์ผ */
41
  .gradio-container {
 
 
 
 
42
  background: linear-gradient(135deg, #EFF6FF 0%, #F5F3FF 100%);
43
- max-width: 1200px !important;
44
- margin: 0 auto !important;
45
- padding: 2rem !important;
46
  }
47
- /* ์ œ๋ชฉ ์Šคํƒ€์ผ */
48
  .title-area {
49
  text-align: center;
50
- margin-bottom: 2rem;
51
  padding: 1rem;
52
  background: white;
53
  border-radius: 1rem;
54
  box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
 
55
  }
56
  .title-area h1 {
57
  background: linear-gradient(90deg, #2563EB 0%, #7C3AED 100%);
@@ -65,81 +76,11 @@ def create_css():
65
  color: #6B7280;
66
  font-size: 1.1rem;
67
  }
68
- /* ์ปดํฌ๋„ŒํŠธ ์Šคํƒ€์ผ๋ง */
69
- .gr-box, .gr-panel {
70
- border: 2px solid #E0E7FF !important;
71
- border-radius: 12px !important;
72
- box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1) !important;
73
- background: white !important;
74
- }
75
- /* ํŒŒ์ผ ์—…๋กœ๋“œ ์˜์—ญ */
76
- .file-upload {
77
- border: 2px dashed #93C5FD !important;
78
- border-radius: 8px !important;
79
- padding: 2rem !important;
80
- background: #F0F9FF !important;
81
- transition: all 0.3s ease;
82
- }
83
- .file-upload:hover {
84
- background: #E0F2FE !important;
85
- border-color: #60A5FA !important;
86
- }
87
- /* ๋ฒ„ํŠผ ์Šคํƒ€์ผ๋ง */
88
- .gr-button.primary-button {
89
- background: linear-gradient(90deg, #2563EB 0%, #7C3AED 100%) !important;
90
- color: white !important;
91
- border: none !important;
92
- border-radius: 8px !important;
93
- padding: 0.75rem 1.5rem !important;
94
- font-weight: bold !important;
95
- transition: opacity 0.2s !important;
96
- }
97
- .gr-button.primary-button:hover {
98
- opacity: 0.9 !important;
99
- }
100
- .gr-button.secondary-button {
101
- background: white !important;
102
- color: #4B5563 !important;
103
- border: 1px solid #D1D5DB !important;
104
- border-radius: 8px !important;
105
- padding: 0.75rem 1.5rem !important;
106
- }
107
- .gr-button.secondary-button:hover {
108
- background: #F9FAFB !important;
109
- }
110
- /* ์Šฌ๋ผ์ด๋” ์Šคํƒ€์ผ๋ง */
111
- .gr-slider {
112
- background: #E0E7FF !important;
113
- }
114
- .gr-slider .gr-slider-handle {
115
- background: #4F46E5 !important;
116
- }
117
- /* ์ฒดํฌ๋ฐ•์Šค ์Šคํƒ€์ผ๋ง */
118
- .gr-checkbox {
119
- border-color: #6366F1 !important;
120
  }
121
- .gr-checkbox:checked {
122
- background-color: #4F46E5 !important;
123
- }
124
- /* ํƒญ ์Šคํƒ€์ผ๋ง */
125
- .gr-tabs {
126
- border-bottom: 2px solid #E0E7FF !important;
127
- }
128
- .gr-tab-button {
129
- color: #6B7280 !important;
130
- padding: 0.75rem 1rem !important;
131
- font-weight: 500 !important;
132
- }
133
- .gr-tab-button.selected {
134
- color: #4F46E5 !important;
135
- border-bottom: 2px solid #4F46E5 !important;
136
- }
137
- /* ๋งˆํฌ๋‹ค์šด ์ถœ๋ ฅ ์˜์—ญ */
138
- .markdown-output {
139
- background: white !important;
140
- border-radius: 8px !important;
141
- padding: 1rem !important;
142
- box-shadow: inset 0 2px 4px rgba(0, 0, 0, 0.05) !important;
143
  }
144
  """
145
 
@@ -149,14 +90,10 @@ def read_fn(path):
149
 
150
  def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, layout_mode, formula_enable, table_enable, language):
151
  os.makedirs(output_dir, exist_ok=True)
152
-
153
  try:
154
  file_name = f"{str(Path(doc_path).stem)}_{time.time()}"
155
  pdf_data = read_fn(doc_path)
156
- if is_ocr:
157
- parse_method = "ocr"
158
- else:
159
- parse_method = "auto"
160
  local_image_dir, local_md_dir = prepare_env(output_dir, file_name, parse_method)
161
  do_parse(
162
  output_dir,
@@ -170,7 +107,7 @@ def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, layout_mode, formula_en
170
  formula_enable=formula_enable,
171
  table_enable=table_enable,
172
  lang=language,
173
- f_dump_orig_pdf=False,
174
  )
175
  return local_md_dir, file_name
176
  except Exception as e:
@@ -202,48 +139,69 @@ def replace_image_with_base64(markdown_text, image_dir_path):
202
  return f"![{relative_path}](data:image/jpeg;base64,{base64_image})"
203
  return re.sub(pattern, replace, markdown_text)
204
 
205
- def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
206
  file_path = to_pdf(file_path)
 
 
207
  if end_pages > 20:
208
  end_pages = 20
 
 
209
  local_md_dir, file_name = parse_pdf(file_path, './output', end_pages - 1, is_ocr,
210
- layout_mode, formula_enable, table_enable, language)
 
 
 
211
  archive_zip_path = os.path.join("./output", compute_sha256(local_md_dir) + ".zip")
212
  zip_archive_success = compress_directory_to_zip(local_md_dir, archive_zip_path)
213
  if zip_archive_success == 0:
214
  logger.info("์••์ถ• ์„ฑ๊ณต")
215
  else:
216
  logger.error("์••์ถ• ์‹คํŒจ")
 
 
 
217
  md_path = os.path.join(local_md_dir, file_name + ".md")
218
  with open(md_path, 'r', encoding='utf-8') as f:
219
  txt_content = f.read()
 
 
 
220
  md_content = replace_image_with_base64(txt_content, local_md_dir)
221
- new_pdf_path = os.path.join(local_md_dir, file_name + "_layout.pdf")
222
- return md_content, txt_content, archive_zip_path, new_pdf_path
223
 
224
- def to_pdf(file_path):
225
- with pymupdf.open(file_path) as f:
226
- if f.is_pdf:
227
- return file_path
228
- else:
229
- pdf_bytes = f.convert_to_pdf()
230
- unique_filename = f"{uuid.uuid4()}.pdf"
231
- tmp_file_path = os.path.join(os.path.dirname(file_path), unique_filename)
232
- with open(tmp_file_path, 'wb') as tmp_pdf_file:
233
- tmp_pdf_file.write(pdf_bytes)
234
- return tmp_file_path
235
 
236
- latex_delimiters = [{"left": "$$", "right": "$$", "display": True},
237
- {"left": '$', "right": '$', "display": False}]
238
 
 
 
 
239
  def init_model():
240
  from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
241
  try:
242
  model_manager = ModelSingleton()
243
  txt_model = model_manager.get_model(False, False)
244
- logger.info(f"txt_model init final")
245
  ocr_model = model_manager.get_model(True, False)
246
- logger.info(f"ocr_model init final")
247
  return 0
248
  except Exception as e:
249
  logger.exception(e)
@@ -253,161 +211,221 @@ model_init = init_model()
253
  logger.info(f"model_init: {model_init}")
254
 
255
  latin_lang = [
256
- 'af', 'az', 'bs', 'cs', 'cy', 'da', 'de', 'es', 'et', 'fr', 'ga', 'hr',
257
- 'hu', 'id', 'is', 'it', 'ku', 'la', 'lt', 'lv', 'mi', 'ms', 'mt', 'nl',
258
- 'no', 'oc', 'pi', 'pl', 'pt', 'ro', 'rs_latin', 'sk', 'sl', 'sq', 'sv',
259
- 'sw', 'tl', 'tr', 'uz', 'vi', 'french', 'german'
260
  ]
261
- arabic_lang = ['ar', 'fa', 'ug', 'ur']
262
- cyrillic_lang = [
263
- 'ru', 'rs_cyrillic', 'be', 'bg', 'uk', 'mn', 'abq', 'ady', 'kbd', 'ava',
264
- 'dar', 'inh', 'che', 'lbe', 'lez', 'tab'
265
- ]
266
- devanagari_lang = [
267
- 'hi', 'mr', 'ne', 'bh', 'mai', 'ang', 'bho', 'mah', 'sck', 'new', 'gom',
268
- 'sa', 'bgc'
269
- ]
270
- other_lang = ['ch', 'en', 'korean', 'japan', 'chinese_cht', 'ta', 'te', 'ka']
271
 
272
  all_lang = ['', 'auto']
273
  all_lang.extend([*other_lang, *latin_lang, *arabic_lang, *cyrillic_lang, *devanagari_lang])
274
 
275
- if __name__ == "__main__":
276
- with gr.Blocks(title="OCR FLEX", css=create_css()) as demo:
277
- # ํƒ€์ดํ‹€ ์˜์—ญ
278
- with gr.Row(elem_classes="title-area"):
279
- gr.HTML("""
280
- <h1>OCR FLEX</h1>
281
- <p>PDF์™€ ์ด๋ฏธ์ง€์—์„œ ํ…์ŠคํŠธ๋ฅผ ๋น ๋ฅด๊ณ  ์ •ํ™•ํ•˜๊ฒŒ ์ถ”์ถœํ•˜์„ธ์š”</p>
282
- """)
283
-
284
- with gr.Row():
285
- # ์™ผ์ชฝ ํŒจ๋„
286
- with gr.Column(variant='panel', scale=5):
287
- file = gr.File(
288
- label="PDF ๋˜๋Š” ์ด๋ฏธ์ง€ ํŒŒ์ผ์„ ์—…๋กœ๋“œํ•˜์„ธ์š”",
289
- file_types=[".pdf", ".png", ".jpeg", ".jpg"],
290
- elem_classes="file-upload"
291
- )
292
-
293
- max_pages = gr.Slider(
294
- 1, 20, 10,
295
- step=1,
296
- label='์ตœ๋Œ€ ๋ณ€ํ™˜ ํŽ˜์ด์ง€ ์ˆ˜',
297
- elem_classes="custom-slider"
298
- )
299
-
300
- with gr.Row():
301
- layout_mode = gr.Dropdown(
302
- ["layoutlmv3", "doclayout_yolo"],
303
- label="๋ ˆ์ด์•„์›ƒ ๋ชจ๋ธ",
304
- value="doclayout_yolo",
305
- elem_classes="custom-dropdown"
306
- )
307
- language = gr.Dropdown(
308
- all_lang,
309
- label="์–ธ์–ด",
310
- value='auto',
311
- elem_classes="custom-dropdown"
312
- )
313
-
314
- with gr.Row():
315
- formula_enable = gr.Checkbox(
316
- label="์ˆ˜์‹ ์ธ์‹ ํ™œ์„ฑํ™”",
317
- value=True,
318
- elem_classes="custom-checkbox"
319
- )
320
- is_ocr = gr.Checkbox(
321
- label="OCR ๊ฐ•์ œ ํ™œ์„ฑํ™”",
322
- value=False,
323
- elem_classes="custom-checkbox"
324
- )
325
- table_enable = gr.Checkbox(
326
- label="ํ‘œ ์ธ์‹ ํ™œ์„ฑํ™”(ํ…Œ์ŠคํŠธ)",
327
- value=True,
328
- elem_classes="custom-checkbox"
329
- )
330
-
331
- with gr.Row():
332
- change_bu = gr.Button(
333
- "๋ณ€ํ™˜",
334
- elem_classes="primary-button"
335
- )
336
- clear_bu = gr.ClearButton(
337
- value="์ดˆ๊ธฐํ™”",
338
- elem_classes="secondary-button"
339
- )
340
-
341
- pdf_show = PDF(
342
- label='PDF ๋ฏธ๋ฆฌ๋ณด๊ธฐ',
343
- interactive=False,
344
- visible=True,
345
- height=800,
346
- elem_classes="pdf-preview"
347
  )
348
-
349
- with gr.Accordion("์˜ˆ์ œ:", open=False):
350
- example_root = os.path.join(os.path.dirname(__file__), "examples")
351
- gr.Examples(
352
- examples=[os.path.join(example_root, _) for _ in os.listdir(example_root) if
353
- _.endswith("pdf")],
354
- inputs=file
355
- )
356
-
357
- # ์˜ค๋ฅธ์ชฝ ํŒจ๋„
358
- with gr.Column(variant='panel', scale=5):
359
- output_file = gr.File(
360
- label="๋ณ€ํ™˜ ๊ฒฐ๊ณผ",
361
- interactive=False,
362
- elem_classes="output-file"
 
 
363
  )
364
-
365
- with gr.Tabs() as tabs:
366
- with gr.Tab("๋งˆํฌ๋‹ค์šด ๋ Œ๋”๋ง"):
367
- md = gr.Markdown(
368
- label="๋งˆํฌ๋‹ค์šด ๋ Œ๋”๋ง",
369
- height=1100,
370
- show_copy_button=True,
371
- latex_delimiters=latex_delimiters,
372
- line_breaks=True,
373
- elem_classes="markdown-output"
374
- )
375
-
376
- with gr.Tab("๋งˆํฌ๋‹ค์šด ํ…์ŠคํŠธ"):
377
- md_text = gr.TextArea(
378
- lines=45,
379
- show_copy_button=True,
380
- elem_classes="markdown-text"
381
- )
382
-
383
- # ์ด๋ฒคํŠธ ํ•ธ๋“ค๋Ÿฌ
384
- file.change(
385
- fn=to_pdf,
386
- inputs=file,
387
- outputs=pdf_show
388
- )
389
-
390
- change_bu.click(
391
- fn=to_markdown,
392
- inputs=[
393
- file,
394
- max_pages,
395
- is_ocr,
396
- layout_mode,
397
- formula_enable,
398
- table_enable,
399
- language
400
- ],
401
- outputs=[
402
- md,
403
- md_text,
404
- output_file,
405
- pdf_show
406
- ],
407
- api_name=False
408
- )
409
-
410
- clear_bu.add([file, md, pdf_show, md_text, output_file, is_ocr])
411
 
412
- # ์•ฑ ์‹คํ–‰
413
- demo.launch(ssr_mode=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8
  import uuid
9
  import pymupdf
10
 
11
+ # =======================================================
12
+ # magic-pdf & MinerU ์„ค์น˜ ๊ด€๋ จ (๊ธฐ์กด ์ฝ”๋“œ ๊ทธ๋Œ€๋กœ ์˜ˆ์‹œ)
13
+ # =======================================================
14
  os.system('pip uninstall -y magic-pdf')
15
  os.system('pip install git+https://github.com/opendatalab/MinerU.git@dev')
 
16
  os.system('wget https://github.com/opendatalab/MinerU/raw/dev/scripts/download_models_hf.py -O download_models_hf.py')
17
  os.system('python download_models_hf.py')
18
 
 
28
  json.dump(data, file, indent=4)
29
 
30
  os.system('cp -r paddleocr /home/user/.paddleocr')
 
31
 
32
  import gradio as gr
33
  from loguru import logger
 
36
  from magic_pdf.libs.hash_utils import compute_sha256
37
  from magic_pdf.tools.common import do_parse, prepare_env
38
 
39
+
40
+ ###########################################
41
+ # 1) UI ์Šคํƒ€์ผ(CSS) + PDF์ฒ˜๋ฆฌ ๊ด€๋ จ ํ•จ์ˆ˜๋“ค
42
+ ###########################################
43
  def create_css():
44
+ """
45
+ ํ™”๋ฉด์„ ๊ฐ€๋“ ์ฑ„์šฐ๊ณ  ์Šคํฌ๋กค ๊ฐ€๋Šฅํ•˜๋„๋ก ์„ค์ •
46
+ """
47
  return """
 
48
  .gradio-container {
49
+ width: 100vw !important;
50
+ min-height: 100vh !important;
51
+ margin: 0 !important;
52
+ padding: 0 !important;
53
  background: linear-gradient(135deg, #EFF6FF 0%, #F5F3FF 100%);
54
+ display: flex;
55
+ flex-direction: column;
56
+ overflow-y: auto !important;
57
  }
 
58
  .title-area {
59
  text-align: center;
60
+ margin: 1rem auto;
61
  padding: 1rem;
62
  background: white;
63
  border-radius: 1rem;
64
  box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
65
+ max-width: 800px;
66
  }
67
  .title-area h1 {
68
  background: linear-gradient(90deg, #2563EB 0%, #7C3AED 100%);
 
76
  color: #6B7280;
77
  font-size: 1.1rem;
78
  }
79
+ .invisible {
80
+ display: none !important;
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
81
  }
82
+ .gr-block, .gr-box {
83
+ padding: 0.5rem !important;
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84
  }
85
  """
86
 
 
90
 
91
  def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, layout_mode, formula_enable, table_enable, language):
92
  os.makedirs(output_dir, exist_ok=True)
 
93
  try:
94
  file_name = f"{str(Path(doc_path).stem)}_{time.time()}"
95
  pdf_data = read_fn(doc_path)
96
+ parse_method = "ocr" if is_ocr else "auto"
 
 
 
97
  local_image_dir, local_md_dir = prepare_env(output_dir, file_name, parse_method)
98
  do_parse(
99
  output_dir,
 
107
  formula_enable=formula_enable,
108
  table_enable=table_enable,
109
  lang=language,
110
+ f_dump_orig_pdf=False
111
  )
112
  return local_md_dir, file_name
113
  except Exception as e:
 
139
  return f"![{relative_path}](data:image/jpeg;base64,{base64_image})"
140
  return re.sub(pattern, replace, markdown_text)
141
 
142
+ def to_pdf(file_path):
143
+ with pymupdf.open(file_path) as f:
144
+ if f.is_pdf:
145
+ return file_path
146
+ else:
147
+ pdf_bytes = f.convert_to_pdf()
148
+ unique_filename = f"{uuid.uuid4()}.pdf"
149
+ tmp_file_path = os.path.join(os.path.dirname(file_path), unique_filename)
150
+ with open(tmp_file_path, 'wb') as tmp_pdf_file:
151
+ tmp_pdf_file.write(pdf_bytes)
152
+ return tmp_file_path
153
+
154
+ def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language, progress=gr.Progress(track_tqdm=False)):
155
+ """
156
+ - PDF ๋ณ€ํ™˜ ๊ณผ์ •์—์„œ 'progress(...)' ๋กœ ์ง„ํ–‰๋ฅ  ์—…๋ฐ์ดํŠธ
157
+ - Gradio ๋ฒ„์ „์ด ๋‚ฎ์•„๋„ 'with progress:' ๋ฅผ ์“ฐ์ง€ ์•Š์œผ๋ฉด __enter__ ์˜ค๋ฅ˜๊ฐ€ ์•ˆ๋œธ
158
+ """
159
+ progress(0, "PDF๋กœ ๋ณ€ํ™˜ ์ค‘...")
160
  file_path = to_pdf(file_path)
161
+ time.sleep(0.5)
162
+
163
  if end_pages > 20:
164
  end_pages = 20
165
+
166
+ progress(20, "๋ฌธ์„œ ํŒŒ์‹ฑ ์ค‘...")
167
  local_md_dir, file_name = parse_pdf(file_path, './output', end_pages - 1, is_ocr,
168
+ layout_mode, formula_enable, table_enable, language)
169
+ time.sleep(0.5)
170
+
171
+ progress(50, "์••์ถ•(zip) ์ƒ์„ฑ ์ค‘...")
172
  archive_zip_path = os.path.join("./output", compute_sha256(local_md_dir) + ".zip")
173
  zip_archive_success = compress_directory_to_zip(local_md_dir, archive_zip_path)
174
  if zip_archive_success == 0:
175
  logger.info("์••์ถ• ์„ฑ๊ณต")
176
  else:
177
  logger.error("์••์ถ• ์‹คํŒจ")
178
+ time.sleep(0.5)
179
+
180
+ progress(70, "๋งˆํฌ๋‹ค์šด ๋กœ๋“œ ์ค‘...")
181
  md_path = os.path.join(local_md_dir, file_name + ".md")
182
  with open(md_path, 'r', encoding='utf-8') as f:
183
  txt_content = f.read()
184
+ time.sleep(0.5)
185
+
186
+ progress(90, "์ด๋ฏธ์ง€(base64) ๋ณ€ํ™˜ ์ค‘...")
187
  md_content = replace_image_with_base64(txt_content, local_md_dir)
188
+ time.sleep(0.5)
 
189
 
190
+ progress(100, "๋ณ€ํ™˜ ์™„๋ฃŒ!")
191
+ return md_content
 
 
 
 
 
 
 
 
 
192
 
 
 
193
 
194
+ ###############################
195
+ # magic_pdf ๋ชจ๋ธ ์ดˆ๊ธฐํ™”
196
+ ###############################
197
  def init_model():
198
  from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
199
  try:
200
  model_manager = ModelSingleton()
201
  txt_model = model_manager.get_model(False, False)
202
+ logger.info("txt_model init final")
203
  ocr_model = model_manager.get_model(True, False)
204
+ logger.info("ocr_model init final")
205
  return 0
206
  except Exception as e:
207
  logger.exception(e)
 
211
  logger.info(f"model_init: {model_init}")
212
 
213
  latin_lang = [
214
+ 'af','az','bs','cs','cy','da','de','es','et','fr','ga','hr','hu','id','is','it','ku',
215
+ 'la','lt','lv','mi','ms','mt','nl','no','oc','pi','pl','pt','ro','rs_latin','sk','sl',
216
+ 'sq','sv','sw','tl','tr','uz','vi','french','german'
 
217
  ]
218
+ arabic_lang = ['ar','fa','ug','ur']
219
+ cyrillic_lang = ['ru','rs_cyrillic','be','bg','uk','mn','abq','ady','kbd','ava','dar','inh','che','lbe','lez','tab']
220
+ devanagari_lang = ['hi','mr','ne','bh','mai','ang','bho','mah','sck','new','gom','sa','bgc']
221
+ other_lang = ['ch','en','korean','japan','chinese_cht','ta','te','ka']
 
 
 
 
 
 
222
 
223
  all_lang = ['', 'auto']
224
  all_lang.extend([*other_lang, *latin_lang, *arabic_lang, *cyrillic_lang, *devanagari_lang])
225
 
226
+
227
+ #################################
228
+ # 2) Gemini (google.generativeai)
229
+ #################################
230
+ import google.generativeai as genai
231
+ from gradio import ChatMessage
232
+ from typing import Iterator
233
+ import time
234
+
235
+ GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
236
+ genai.configure(api_key=GEMINI_API_KEY)
237
+
238
+ model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp-1219")
239
+
240
+ def format_chat_history(messages: list) -> list:
241
+ """
242
+ Gemini๊ฐ€ ์ดํ•ดํ•  ์ˆ˜ ์žˆ๋Š” ํ˜•์‹ (role, content)
243
+ """
244
+ formatted_history = []
245
+ for message in messages:
246
+ if not (message.role == "assistant" and hasattr(message, "metadata")):
247
+ formatted_history.append({
248
+ "role": "user" if message.role == "user" else "assistant",
249
+ "parts": [message.content]
250
+ })
251
+ return formatted_history
252
+
253
+ def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
254
+ """
255
+ Gemini ์‘๋‹ต ์ŠคํŠธ๋ฆฌ๋ฐ: user_message๊ฐ€ ๋น„์–ด์žˆ์œผ๋ฉด ๊ธฐ๋ณธ ๋ฌธ๊ตฌ๋กœ ๋Œ€์ฒด
256
+ """
257
+ # ๋นˆ ๋ฌธ์ž์—ด์ด๋ฉด ๊ธฐ๋ณธ ๋ฌธ๊ตฌ๋กœ ๊ต์ฒด(์˜ค๋ฅ˜ ๋ฐฉ์ง€)
258
+ if not user_message.strip():
259
+ user_message = "โ€ฆ(No content from user)โ€ฆ"
260
+
261
+ try:
262
+ print(f"\n=== [Gemini] New Request ===\nUser message: '{user_message}'")
263
+
264
+ chat_history = format_chat_history(messages)
265
+ chat = model.start_chat(history=chat_history)
266
+ response = chat.send_message(user_message, stream=True)
267
+
268
+ thought_buffer = ""
269
+ response_buffer = ""
270
+ thinking_complete = False
271
+
272
+ # "Thinking" ๋ฉ”์‹œ์ง€ ์ถ”๊ฐ€
273
+ messages.append(
274
+ ChatMessage(
275
+ role="assistant",
276
+ content="",
277
+ metadata={"title": "โš™๏ธ Thinking: *The thoughts produced by the model are experimental"}
278
+ )
279
+ )
280
+ yield convert_chat_messages_to_gradio_format(messages)
281
+
282
+ for chunk in response:
283
+ parts = chunk.candidates[0].content.parts
284
+ current_chunk = parts[0].text
285
+
286
+ if len(parts) == 2 and not thinking_complete:
287
+ # Complete thought
288
+ thought_buffer += current_chunk
289
+ messages[-1] = ChatMessage(
290
+ role="assistant",
291
+ content=thought_buffer,
292
+ metadata={"title": "โš™๏ธ Thinking: *The thoughts produced by the model are experimental"}
 
 
 
 
 
293
  )
294
+ yield convert_chat_messages_to_gradio_format(messages)
295
+
296
+ # Start final response
297
+ response_buffer = parts[1].text
298
+ messages.append(ChatMessage(role="assistant", content=response_buffer))
299
+ thinking_complete = True
300
+ elif thinking_complete:
301
+ # Response ongoing
302
+ response_buffer += current_chunk
303
+ messages[-1] = ChatMessage(role="assistant", content=response_buffer)
304
+ else:
305
+ # Still thinking
306
+ thought_buffer += current_chunk
307
+ messages[-1] = ChatMessage(
308
+ role="assistant",
309
+ content=thought_buffer,
310
+ metadata={"title": "โš™๏ธ Thinking: *The thoughts produced by the model are experimental"}
311
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
312
 
313
+ yield convert_chat_messages_to_gradio_format(messages)
314
+
315
+ print(f"\n=== [Gemini] Final Response ===\n{response_buffer}")
316
+
317
+ except Exception as e:
318
+ print(f"\n=== [Gemini] Error ===\n{str(e)}")
319
+ messages.append(ChatMessage(role="assistant", content=f"I encountered an error: {str(e)}"))
320
+ yield convert_chat_messages_to_gradio_format(messages)
321
+
322
+ def convert_chat_messages_to_gradio_format(messages):
323
+ """
324
+ ChatMessage list -> [ (์œ ์ €๋ฐœํ™”, ๋ด‡์‘๋‹ต), (...), ... ]
325
+ """
326
+ gradio_chat = []
327
+ user_text, assistant_text = None, None
328
+
329
+ for msg in messages:
330
+ if msg.role == "user":
331
+ # ์ด์ „ ํ„ด ์ €์žฅ
332
+ if user_text is not None or assistant_text is not None:
333
+ gradio_chat.append((user_text or "", assistant_text or ""))
334
+ user_text = msg.content
335
+ assistant_text = None
336
+ else:
337
+ # assistant
338
+ if user_text is None:
339
+ user_text = ""
340
+ if assistant_text is None:
341
+ assistant_text = msg.content
342
+ else:
343
+ assistant_text += msg.content # ์ŠคํŠธ๋ฆฌ๋ฐ ์‹œ ๋ˆ„์ 
344
+
345
+ # ๋งˆ์ง€๋ง‰ ํ„ด
346
+ if user_text is not None or assistant_text is not None:
347
+ gradio_chat.append((user_text or "", assistant_text or ""))
348
+
349
+ return gradio_chat
350
+
351
+ def user_message(msg: str, history: list, doc_text: str) -> tuple[str, list]:
352
+ """
353
+ doc_text(๋งˆํฌ๋‹ค์šด) ์ฐธ๊ณ  ๋ฌธ๊ตฌ๋ฅผ ์ž๋™ ์‚ฝ์ž…
354
+ """
355
+ if doc_text.strip():
356
+ user_query = f"๋‹ค์Œ ๋ฌธ์„œ๋ฅผ ์ฐธ๊ณ ํ•˜์—ฌ ๋‹ต๋ณ€:\n\n{doc_text}\n\n์งˆ๋ฌธ: {msg}"
357
+ else:
358
+ user_query = msg
359
+
360
+ history.append(ChatMessage(role="user", content=user_query))
361
+ return "", history
362
+
363
+
364
+ ################################
365
+ # 3) ํ†ตํ•ฉ Gradio ์•ฑ ๊ตฌ์„ฑ & ์‹คํ–‰
366
+ ################################
367
+ with gr.Blocks(title="OCR FLEX + Gemini Chat", css=create_css()) as demo:
368
+ gr.HTML("""
369
+ <div class="title-area">
370
+ <h1>OCR FLEX + Gemini Chat</h1>
371
+ <p>PDF/์ด๋ฏธ์ง€ -> ํ…์ŠคํŠธ(๋งˆํฌ๋‹ค์šด) ๋ณ€ํ™˜ ํ›„, Gemini LLM ๋Œ€ํ™”</p>
372
+ </div>
373
+ """)
374
+
375
+ md_state = gr.State("")
376
+ chat_history = gr.State([])
377
+
378
+ with gr.Row():
379
+ file = gr.File(label="PDF/์ด๋ฏธ์ง€ ์—…๋กœ๋“œ", file_types=[".pdf", ".png", ".jpeg", ".jpg"], interactive=True)
380
+ convert_btn = gr.Button("๋ณ€ํ™˜ํ•˜๊ธฐ")
381
+
382
+ # ์ˆจ๊ธด ์ปดํฌ๋„ŒํŠธ๋“ค
383
+ max_pages = gr.Slider(1, 20, 10, visible=False, elem_classes="invisible")
384
+ layout_mode = gr.Dropdown(["layoutlmv3","doclayout_yolo"], value="doclayout_yolo", visible=False, elem_classes="invisible")
385
+ language = gr.Dropdown(all_lang, value='auto', visible=False, elem_classes="invisible")
386
+ formula_enable = gr.Checkbox(value=True, visible=False, elem_classes="invisible")
387
+ is_ocr = gr.Checkbox(value=False, visible=False, elem_classes="invisible")
388
+ table_enable = gr.Checkbox(value=True, visible=False, elem_classes="invisible")
389
+
390
+ # ๋ณ€ํ™˜ ํด๋ฆญ -> to_markdown (progress)
391
+ convert_btn.click(
392
+ fn=to_markdown,
393
+ inputs=[file, max_pages, is_ocr, layout_mode, formula_enable, table_enable, language],
394
+ outputs=md_state,
395
+ show_progress=True # ํ”„๋กœ๊ทธ๋ ˆ์Šค๋ฐ”+๋กœ๋”ฉ ํ‘œ์‹œ
396
+ )
397
+
398
+ # Gemini Chat
399
+ gr.Markdown("## Gemini 2.0 Flash (Thinking) Chat")
400
+ chatbot = gr.Chatbot(height=600)
401
+ with gr.Row():
402
+ chat_input = gr.Textbox(lines=1, placeholder="์งˆ๋ฌธ์„ ์ž…๋ ฅํ•˜์„ธ์š”...")
403
+ clear_btn = gr.Button("๋Œ€ํ™” ์ดˆ๊ธฐํ™”")
404
+
405
+ # ํ”„๋กฌํ”„ํŠธ ์ „์†ก -> user_message -> stream_gemini_response
406
+ chat_input.submit(
407
+ fn=user_message,
408
+ inputs=[chat_input, chat_history, md_state],
409
+ outputs=[chat_input, chat_history]
410
+ ).then(
411
+ fn=stream_gemini_response,
412
+ inputs=[chat_input, chat_history],
413
+ outputs=chatbot
414
+ )
415
+
416
+ def clear_states():
417
+ return [], ""
418
+
419
+ clear_btn.click(
420
+ fn=clear_states,
421
+ inputs=[],
422
+ outputs=[chat_history, md_state]
423
+ ).then(
424
+ fn=lambda: [],
425
+ inputs=[],
426
+ outputs=chatbot
427
+ )
428
+
429
+
430
+ if __name__ == "__main__":
431
+ demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)