pnnbao-ump commited on
Commit
2945c35
·
1 Parent(s): 8f62483

first init

Browse files
README.md CHANGED
@@ -1,14 +1,54 @@
1
  ---
2
- title: MedCrab
3
- emoji: 🏃
4
- colorFrom: yellow
5
- colorTo: purple
6
  sdk: gradio
7
- sdk_version: 6.0.1
8
  app_file: app.py
9
- pinned: false
10
  license: cc-by-nc-4.0
11
  short_description: Medical PDF Translator
12
  ---
13
 
14
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ title: MedCrab Translation
3
+ emoji: 🦀
4
+ colorFrom: purple
5
+ colorTo: blue
6
  sdk: gradio
7
+ sdk_version: 5.0.0
8
  app_file: app.py
9
+ pinned: true
10
  license: cc-by-nc-4.0
11
  short_description: Medical PDF Translator
12
  ---
13
 
14
+ # 🦀 MedCrab Translation
15
+
16
+ Ứng dụng quét OCR tài liệu y khoa và dịch trực tiếp sang tiếng Việt với hiệu ứng streaming.
17
+
18
+ ## Tính năng
19
+
20
+ - 📄 Hỗ trợ PDF và hình ảnh
21
+ - 🔍 OCR chính xác với DeepSeek-OCR
22
+ - 🦀 Dịch y khoa chuyên sâu với MedCrab
23
+ - ⚡ Streaming real-time
24
+ - 🎨 Giao diện thân thiện
25
+
26
+ ## Sử dụng
27
+
28
+ 1. Tải lên file PDF hoặc hình ảnh y khoa
29
+ 2. Chọn số trang (nếu là PDF)
30
+ 3. Chọn chế độ OCR
31
+ 4. Nhấn "Quét OCR + Dịch tiếng Việt"
32
+
33
+ ## Chế độ OCR
34
+
35
+ - **Crab**: Chế độ cân bằng (khuyên dùng)
36
+ - **Base**: Chế độ nhanh
37
+
38
+ ## Yêu cầu GPU
39
+
40
+ Space này cần GPU để chạy. Hugging Face cung cấp GPU miễn phí với giới hạn thời gian sử dụng.
41
+
42
+ ## License
43
+
44
+ Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0)
45
+
46
+ This work is licensed under CC BY-NC 4.0. You are free to:
47
+ - Share: copy and redistribute the material
48
+ - Adapt: remix, transform, and build upon the material
49
+
50
+ Under the following terms:
51
+ - Attribution: You must give appropriate credit
52
+ - NonCommercial: You may not use the material for commercial purposes
53
+
54
+ See: https://creativecommons.org/licenses/by-nc/4.0/
app.py ADDED
@@ -0,0 +1,471 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModel, AutoTokenizer
3
+ from medcrab import MedCrabTranslator
4
+ import torch
5
+ import os
6
+ import sys
7
+ import tempfile
8
+ import shutil
9
+ from PIL import Image, ImageOps
10
+ import fitz
11
+ import re
12
+ import time
13
+ from threading import Thread
14
+ from queue import Queue
15
+ from io import StringIO, BytesIO
16
+ import spaces
17
+
18
+ # ==================== DEEPSEEK OCR SETUP ====================
19
+ OCR_MODEL_NAME = 'deepseek-ai/DeepSeek-OCR'
20
+
21
+ print("🔄 Loading OCR model...")
22
+ ocr_tokenizer = AutoTokenizer.from_pretrained(OCR_MODEL_NAME, trust_remote_code=True)
23
+
24
+ try:
25
+ ocr_model = AutoModel.from_pretrained(
26
+ OCR_MODEL_NAME,
27
+ attn_implementation='flash_attention_2',
28
+ torch_dtype=torch.bfloat16,
29
+ trust_remote_code=True,
30
+ use_safetensors=True
31
+ )
32
+ print("✅ Using Flash Attention 2")
33
+ except (ImportError, ValueError):
34
+ print("⚠️ Flash Attention 2 not available, using eager attention")
35
+ ocr_model = AutoModel.from_pretrained(
36
+ OCR_MODEL_NAME,
37
+ attn_implementation='eager',
38
+ torch_dtype=torch.bfloat16,
39
+ trust_remote_code=True,
40
+ use_safetensors=True
41
+ )
42
+
43
+ ocr_model = ocr_model.eval()
44
+
45
+ MODEL_CONFIGS = {
46
+ "Crab": {"base_size": 1024, "image_size": 640, "crop_mode": True},
47
+ "Base": {"base_size": 1024, "image_size": 1024, "crop_mode": False},
48
+ }
49
+
50
+ # ==================== MEDCRAB TRANSLATOR SETUP ====================
51
+ print("🦀 Loading MedCrab translator...")
52
+ translator = None
53
+
54
+ def init_translator():
55
+ global translator
56
+ if translator is None:
57
+ device = "cuda" if torch.cuda.is_available() else "cpu"
58
+ translator = MedCrabTranslator(device=device)
59
+ print(f"✅ MedCrab translator loaded on {device}")
60
+
61
+ # ==================== TEXT CLEANING FUNCTIONS ====================
62
+ def clean_mathrm(text):
63
+ """Chuyển đổi LaTeX sang HTML với subscript/superscript chỉ trong môi trường toán học"""
64
+ if not text:
65
+ return ""
66
+
67
+ def process_math_block(match):
68
+ math_content = match.group(1)
69
+ math_content = re.sub(r'\\mathrm\{([^}]*)\}', r'\1', math_content)
70
+ math_content = re.sub(r'\^\{([^}]+)\}', r'<sup>\1</sup>', math_content)
71
+ math_content = re.sub(r'\^([A-Za-z0-9+\-]+)', r'<sup>\1</sup>', math_content)
72
+ math_content = re.sub(r'_\{([^}]+)\}', r'<sub>\1</sub>', math_content)
73
+ math_content = re.sub(r'_([A-Za-z0-9+\-]+)', r'<sub>\1</sub>', math_content)
74
+
75
+ replacements = {
76
+ r'\times': '×', r'\pm': '±', r'\div': '÷', r'\cdot': '·',
77
+ r'\approx': '≈', r'\leq': '≤', r'\geq': '≥', r'\neq': '≠',
78
+ r'\rightarrow': '→', r'\leftarrow': '←',
79
+ r'\Rightarrow': '⇒', r'\Leftarrow': '⇐',
80
+ }
81
+ for latex_cmd, unicode_char in replacements.items():
82
+ math_content = math_content.replace(latex_cmd, unicode_char)
83
+
84
+ return math_content
85
+
86
+ text = re.sub(r'\\\((.+?)\\\)', process_math_block, text, flags=re.DOTALL)
87
+
88
+ def process_bracket_block(m):
89
+ class FakeMatch:
90
+ def __init__(self, content):
91
+ self.content = content
92
+ def group(self, n):
93
+ return self.content
94
+ content = process_math_block(FakeMatch(m.group(1)))
95
+ return '[' + content + ']'
96
+
97
+ text = re.sub(r'\\\[(.+?)\\\]', process_bracket_block, text, flags=re.DOTALL)
98
+ text = re.sub(r'\\mathrm\{([^}]*)\}', r'\1', text)
99
+ text = text.replace(r'\%', '%')
100
+
101
+ lines = text.split('\n')
102
+ cleaned_lines = [re.sub(r'[ \t]+', ' ', line).strip() for line in lines]
103
+ text = '\n'.join(cleaned_lines)
104
+
105
+ return text.strip()
106
+
107
+ def clean_output(text, include_images=False, remove_labels=False):
108
+ if not text:
109
+ return ""
110
+ pattern = r'(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)'
111
+ matches = re.findall(pattern, text, re.DOTALL)
112
+ img_num = 0
113
+
114
+ for match in matches:
115
+ if '<|ref|>image<|/ref|>' in match[0]:
116
+ if include_images:
117
+ text = text.replace(match[0], f'\n\n**[Figure {img_num + 1}]**\n\n', 1)
118
+ img_num += 1
119
+ else:
120
+ text = text.replace(match[0], '', 1)
121
+ else:
122
+ if remove_labels:
123
+ text = text.replace(match[0], '', 1)
124
+ else:
125
+ text = text.replace(match[0], match[1], 1)
126
+
127
+ text = clean_mathrm(text)
128
+ return text.strip()
129
+
130
+ # ==================== OCR FUNCTIONS ====================
131
+ @spaces.GPU
132
+ def ocr_process_image(image, mode="Crab"):
133
+ if image is None:
134
+ return "Error: Upload image"
135
+
136
+ if image.mode in ('RGBA', 'LA', 'P'):
137
+ image = image.convert('RGB')
138
+ image = ImageOps.exif_transpose(image)
139
+
140
+ config = MODEL_CONFIGS[mode]
141
+ prompt = "<image>\n<|grounding|>Convert the document to markdown."
142
+
143
+ tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg')
144
+ image.save(tmp.name, 'JPEG', quality=95)
145
+ tmp.close()
146
+ out_dir = tempfile.mkdtemp()
147
+
148
+ stdout = sys.stdout
149
+ sys.stdout = StringIO()
150
+
151
+ ocr_model.infer(
152
+ tokenizer=ocr_tokenizer,
153
+ prompt=prompt,
154
+ image_file=tmp.name,
155
+ output_path=out_dir,
156
+ base_size=config["base_size"],
157
+ image_size=config["image_size"],
158
+ crop_mode=config["crop_mode"]
159
+ )
160
+
161
+ result = '\n'.join([l for l in sys.stdout.getvalue().split('\n')
162
+ if not any(s in l for s in ['image:', 'other:', 'PATCHES', '====', 'BASE:', '%|', 'torch.Size'])]).strip()
163
+ sys.stdout = stdout
164
+
165
+ os.unlink(tmp.name)
166
+ shutil.rmtree(out_dir, ignore_errors=True)
167
+
168
+ if not result:
169
+ return "No text detected"
170
+
171
+ markdown = clean_output(result, True, True)
172
+ return markdown
173
+
174
+ def ocr_process_pdf(path, mode, page_num):
175
+ doc = fitz.open(path)
176
+ total_pages = len(doc)
177
+ if page_num < 1 or page_num > total_pages:
178
+ doc.close()
179
+ return f"Invalid page number. PDF has {total_pages} pages."
180
+
181
+ page = doc.load_page(page_num - 1)
182
+ pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72), alpha=False)
183
+ img = Image.open(BytesIO(pix.tobytes("png")))
184
+ doc.close()
185
+
186
+ return ocr_process_image(img, mode)
187
+
188
+ def ocr_process_file(path, mode, page_num):
189
+ if not path:
190
+ return "Error: Upload file"
191
+ if path.lower().endswith('.pdf'):
192
+ return ocr_process_pdf(path, mode, page_num)
193
+ else:
194
+ return ocr_process_image(Image.open(path), mode)
195
+
196
+ # ==================== TRANSLATION FUNCTIONS ====================
197
+ def split_by_sentences(text: str, max_words: int = 100):
198
+ def count_words(t):
199
+ return len(t.strip().split())
200
+
201
+ chunks = []
202
+ lines = text.split('\n')
203
+
204
+ i = 0
205
+ while i < len(lines):
206
+ line = lines[i]
207
+
208
+ empty_count = 0
209
+ if not line.strip():
210
+ while i < len(lines) and not lines[i].strip():
211
+ empty_count += 1
212
+ i += 1
213
+
214
+ if chunks:
215
+ prev_text, prev_newlines = chunks[-1]
216
+ chunks[-1] = (prev_text, prev_newlines + empty_count)
217
+ continue
218
+
219
+ line = line.strip()
220
+ is_last_line = (i == len(lines) - 1)
221
+
222
+ if count_words(line) <= max_words:
223
+ chunks.append((line, 0 if is_last_line else 1))
224
+ i += 1
225
+ continue
226
+
227
+ sentences = re.split(r'(?<=[.!?])\s+', line)
228
+ current_chunk = ""
229
+ current_words = 0
230
+
231
+ for sent_idx, sentence in enumerate(sentences):
232
+ sentence = sentence.strip()
233
+ if not sentence:
234
+ continue
235
+
236
+ sentence_words = count_words(sentence)
237
+
238
+ if sentence_words > max_words:
239
+ if current_chunk:
240
+ chunks.append((current_chunk.strip(), 0))
241
+ current_chunk = ""
242
+ current_words = 0
243
+
244
+ sub_parts = re.split(r',\s*', sentence)
245
+ temp_chunk = ""
246
+ temp_words = 0
247
+
248
+ for part in sub_parts:
249
+ part_words = count_words(part)
250
+ if temp_words + part_words > max_words and temp_chunk:
251
+ chunks.append((temp_chunk.strip(), 0))
252
+ temp_chunk = part
253
+ temp_words = part_words
254
+ else:
255
+ if temp_chunk:
256
+ temp_chunk += ", " + part
257
+ else:
258
+ temp_chunk = part
259
+ temp_words += part_words
260
+
261
+ if temp_chunk.strip():
262
+ current_chunk = temp_chunk.strip()
263
+ current_words = temp_words
264
+
265
+ elif current_words + sentence_words <= max_words:
266
+ if current_chunk:
267
+ current_chunk += " " + sentence
268
+ else:
269
+ current_chunk = sentence
270
+ current_words += sentence_words
271
+ else:
272
+ chunks.append((current_chunk.strip(), 0))
273
+ current_chunk = sentence
274
+ current_words = sentence_words
275
+
276
+ if current_chunk.strip():
277
+ chunks.append((current_chunk.strip(), 0 if is_last_line else 1))
278
+
279
+ i += 1
280
+
281
+ return chunks
282
+
283
+ @spaces.GPU
284
+ def translate_chunk(chunk_text):
285
+ init_translator()
286
+ return translator.translate(chunk_text, max_new_tokens=2048).strip()
287
+
288
+ def streaming_translate(text: str):
289
+ if not text or not text.strip():
290
+ yield '<div style="padding:20px; color:#ff6b6b;">⚠️ Vui lòng nhập văn bản tiếng Anh để dịch.</div>'
291
+ return
292
+
293
+ chunks = split_by_sentences(text, max_words=100)
294
+ accumulated = ""
295
+
296
+ for i, (chunk_text, newline_count) in enumerate(chunks):
297
+ try:
298
+ translated = translate_chunk(chunk_text)
299
+
300
+ if accumulated and not accumulated.endswith('\n'):
301
+ accumulated += " " + translated
302
+ else:
303
+ accumulated += translated
304
+
305
+ chunk_start = len(accumulated) - len(translated)
306
+ for j in range(len(translated)):
307
+ current_display = accumulated[:chunk_start + j + 1]
308
+ html_output = f'<div style="padding:20px; line-height:1.8; font-size:15px; white-space:pre-wrap; font-family:Arial,sans-serif;">{current_display}</div>'
309
+ yield html_output
310
+ time.sleep(0.015)
311
+
312
+ if newline_count > 0:
313
+ actual_newlines = min(newline_count, 2)
314
+ accumulated += "\n" * actual_newlines
315
+ html_output = f'<div style="padding:20px; line-height:1.8; font-size:15px; white-space:pre-wrap; font-family:Arial,sans-serif;">{accumulated}</div>'
316
+ yield html_output
317
+
318
+ except Exception as e:
319
+ yield f'<div style="padding:20px; color:#ff6b6b;">❌ Lỗi dịch chunk {i+1}: {str(e)}</div>'
320
+ return
321
+
322
+ # ==================== UI HELPER FUNCTIONS ====================
323
+ def load_image(file_path, page_num_str="1"):
324
+ if not file_path:
325
+ return None
326
+ try:
327
+ try:
328
+ page_num = int(page_num_str)
329
+ except (ValueError, TypeError):
330
+ page_num = 1
331
+
332
+ if file_path.lower().endswith('.pdf'):
333
+ doc = fitz.open(file_path)
334
+ page_idx = max(0, min(page_num - 1, len(doc) - 1))
335
+ page = doc.load_page(page_idx)
336
+ pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72), alpha=False)
337
+ img = Image.open(BytesIO(pix.tobytes("png")))
338
+ doc.close()
339
+ return img
340
+ else:
341
+ return Image.open(file_path)
342
+ except Exception as e:
343
+ print(f"Error loading image: {e}")
344
+ return None
345
+
346
+ def get_pdf_page_count(file_path):
347
+ if not file_path or not file_path.lower().endswith('.pdf'):
348
+ return 1
349
+ try:
350
+ doc = fitz.open(file_path)
351
+ count = len(doc)
352
+ doc.close()
353
+ return count
354
+ except Exception as e:
355
+ print(f"Error reading PDF page count: {e}")
356
+ return 1
357
+
358
+ def update_page_info(file_path):
359
+ if not file_path:
360
+ return gr.update(label="Số trang (chỉ dùng cho PDF, mặc định: 1)")
361
+ if file_path.lower().endswith('.pdf'):
362
+ page_count = get_pdf_page_count(file_path)
363
+ return gr.update(
364
+ label=f"Số trang (PDF có {page_count} trang, nhập 1-{page_count})",
365
+ value="1"
366
+ )
367
+ return gr.update(
368
+ label="Số trang (chỉ dùng cho PDF, mặc định: 1)",
369
+ value="1"
370
+ )
371
+
372
+ # ==================== COMBINED OCR + TRANSLATION ====================
373
+ def ocr_and_translate_streaming(file_path, mode, page_num_str):
374
+ """Hàm kết hợp: OCR trước, sau đó dịch streaming"""
375
+ if not file_path:
376
+ yield '<div style="padding:20px; color:#ff6b6b;">⚠️ Vui lòng tải file lên trước!</div>'
377
+ return
378
+
379
+ yield '<div style="padding:20px; color:#4CAF50;">🔍 Đang quét OCR...</div>'
380
+ try:
381
+ try:
382
+ page_num = int(page_num_str)
383
+ except (ValueError, TypeError):
384
+ page_num = 1
385
+
386
+ markdown = ocr_process_file(file_path, mode, page_num)
387
+
388
+ if not markdown or markdown.startswith("Error") or markdown.startswith("Invalid"):
389
+ yield f'<div style="padding:20px; color:#ff6b6b;">❌ Lỗi OCR: {markdown}</div>'
390
+ return
391
+
392
+ except Exception as e:
393
+ yield f'<div style="padding:20px; color:#ff6b6b;">❌ Lỗi OCR: {str(e)}</div>'
394
+ return
395
+
396
+ yield '<div style="padding:20px; color:#2196F3;">🦀 Đang dịch...</div>'
397
+ time.sleep(0.5)
398
+
399
+ try:
400
+ yield from streaming_translate(markdown)
401
+ except Exception as e:
402
+ yield f'<div style="padding:20px; color:#ff6b6b;">❌ Lỗi dịch: {str(e)}</div>'
403
+
404
+ # ==================== GRADIO INTERFACE ====================
405
+ css = """
406
+ footer { visibility: hidden }
407
+ .main-title {
408
+ background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
409
+ color: white;
410
+ padding: 15px;
411
+ border-radius: 10px;
412
+ text-align: center;
413
+ margin-bottom: 20px;
414
+ }
415
+ """
416
+
417
+ with gr.Blocks(theme=gr.themes.Soft(), css=css, title="OCR + Translation") as demo:
418
+
419
+ gr.Markdown("""
420
+ <div class="main-title">
421
+ <h1>🦀 MedCrab Translation</h1>
422
+ <p><b>Quét PDF Y khoa → Dịch trực tiếp sang tiếng Việt (Streaming)</b></p>
423
+ </div>
424
+ """)
425
+
426
+ with gr.Row():
427
+ with gr.Column(scale=1):
428
+ gr.Markdown("### 📤 Tải file lên")
429
+ file_in = gr.File(label="PDF hoặc Hình ảnh", file_types=["image", ".pdf"], type="filepath")
430
+ input_img = gr.Image(label="Xem trước", type="pil", height=300)
431
+
432
+ page_input = gr.Textbox(
433
+ label="Số trang (chỉ dùng cho PDF, mặc định: 1)",
434
+ value="1",
435
+ placeholder="Nhập số trang..."
436
+ )
437
+ mode = gr.Dropdown(list(MODEL_CONFIGS.keys()), value="Crab", label="Chế độ OCR")
438
+
439
+ gr.Markdown("### 🦀 Quét và Dịch")
440
+ process_btn = gr.Button("🚀 Quét OCR + Dịch tiếng Việt", variant="primary", size="lg")
441
+
442
+ with gr.Column(scale=2):
443
+ gr.Markdown("### 📄 Kết quả dịch tiếng Việt (Streaming)")
444
+ translation_output = gr.HTML(label="", value="")
445
+
446
+ with gr.Accordion("ℹ️ Hướng dẫn sử dụng", open=False):
447
+ gr.Markdown("""
448
+ **Quy trình đơn giản:**
449
+ 1. 📤 Tải lên file PDF hoặc hình ảnh y khoa
450
+ 2. 🚀 Nhấn nút "Quét OCR + Dịch tiếng Việt"
451
+
452
+ **Chế độ OCR:**
453
+ - **Crab**: 1024 base + 640 tiles (Tốt nhất, cân bằng)
454
+ - **Base**: 1024×1024 (Nhanh hơn)
455
+
456
+ **Lưu ý:** Space này sử dụng GPU miễn phí của Hugging Face, có thể mất vài giây để khởi động.
457
+ """)
458
+
459
+ file_in.change(load_image, [file_in, page_input], [input_img])
460
+ file_in.change(update_page_info, [file_in], [page_input])
461
+ page_input.change(load_image, [file_in, page_input], [input_img])
462
+
463
+ process_btn.click(
464
+ ocr_and_translate_streaming,
465
+ [file_in, mode, page_input],
466
+ [translation_output]
467
+ )
468
+
469
+ if __name__ == "__main__":
470
+ print("🚀 Starting MedCrab Translation on Hugging Face Spaces...")
471
+ demo.queue(max_size=20).launch()
medcrab/__init__.py ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ from .medcrab import MedCrabTranslator
2
+
3
+ __all__ = ["MedCrabTranslator"]
medcrab/__pycache__/__init__.cpython-312.pyc ADDED
Binary file (207 Bytes). View file
 
medcrab/__pycache__/medcrab.cpython-312.pyc ADDED
Binary file (6.88 kB). View file
 
medcrab/medcrab.py ADDED
@@ -0,0 +1,150 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ import re
4
+
5
+ class MedCrabTranslator:
6
+ """
7
+ Translator class cho MedCrab-1.5B, dịch văn bản y khoa từ tiếng Anh sang tiếng Việt.
8
+ """
9
+
10
+ def __init__(self, model_name: str = "pnnbao-ump/MedCrab-1.5b", device: str = "cuda"):
11
+ self.model_name = model_name
12
+ self.device = device
13
+
14
+ # Load tokenizer và model
15
+ print(f"Loading model {model_name} on {device} ...")
16
+ self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
17
+ self.model = AutoModelForCausalLM.from_pretrained(
18
+ model_name,
19
+ device_map="auto" if device == "cuda" else None,
20
+ torch_dtype=torch.bfloat16
21
+ )
22
+ self.model.eval()
23
+
24
+ def _build_prompt(self, text: str) -> str:
25
+ """Tạo prompt chuẩn cho MedCrab"""
26
+ return (
27
+ f"<|user|>: Translate the following medical text from English to Vietnamese:"
28
+ f"<|ENGLISH_TEXT_START|>{text}<|ENGLISH_TEXT_END|>\n"
29
+ "<|assistant|>:<|VIETNAMESE_TEXT_START|>"
30
+ )
31
+
32
+ def translate(self, text: str, max_new_tokens: int = 2048) -> str:
33
+ """
34
+ Dịch văn bản tiếng Anh sang tiếng Việt.
35
+ Args:
36
+ text: Văn bản tiếng Anh
37
+ max_new_tokens: số token tối đa sinh ra
38
+ Returns:
39
+ Văn bản tiếng Việt
40
+ """
41
+ prompt = self._build_prompt(text)
42
+ inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
43
+
44
+ # Lấy EOS token, fallback nếu không có
45
+ eos_id = self.tokenizer.get_vocab().get("<|VIETNAMESE_TEXT_END|>", self.tokenizer.eos_token_id)
46
+
47
+ # Generate với autocast
48
+ with torch.no_grad():
49
+ output_tokens = self.model.generate(
50
+ **inputs,
51
+ max_new_tokens=max_new_tokens,
52
+ eos_token_id=eos_id,
53
+ do_sample=False,
54
+ use_cache=True,
55
+ )
56
+
57
+ # Bỏ phần prompt
58
+ input_len = inputs["input_ids"].shape[-1]
59
+ decoded = self.tokenizer.decode(output_tokens[0, input_len:], skip_special_tokens=True)
60
+
61
+ # Tách token end nếu có
62
+ if "<|VIETNAMESE_TEXT_END|>" in decoded:
63
+ decoded = decoded.split("<|VIETNAMESE_TEXT_END|>")[0]
64
+
65
+ # Optional: post-processing spacing thuật ngữ y khoa
66
+ decoded = self._post_process(decoded)
67
+
68
+ return decoded.strip()
69
+
70
+ def _post_process(self, text: str) -> str:
71
+ """
72
+ Sửa spacing cho các thuật ngữ y khoa bị dính.
73
+ Ví dụ: 'phối tửthụ thể' -> 'phối tử – thụ thể'
74
+ """
75
+ text = text.replace("phối tửthụ thể", "phối tử – thụ thể")
76
+ text = text.replace("miễn dịchchuyển hóa", "miễn dịch – chuyển hóa")
77
+ text = text.replace("oxy hóakhử", "oxy hóa - khử")
78
+ # có thể thêm các quy tắc khác
79
+ return text
80
+
81
+ def _split_into_chunks(self, text: str, max_words: int = 150) -> list[tuple[str, str]]:
82
+ """
83
+ Chia văn bản dài thành các đoạn nhỏ, GHI NHỚ dấu phân cách gốc.
84
+ Return: list[(chunk_text, separator)]
85
+ """
86
+ # Tách câu và GIỮ LẠI dấu phân cách
87
+ pattern = r'([.!?])\n+'
88
+ parts = re.split(pattern, text)
89
+
90
+ sentences = []
91
+ i = 0
92
+ while i < len(parts):
93
+ if i + 1 < len(parts) and parts[i+1] and parts[i+1] in '.!?':
94
+ # Ghép câu với dấu câu của nó
95
+ sentences.append((parts[i] + parts[i+1], parts[i+1]))
96
+ i += 2
97
+ elif parts[i].strip():
98
+ # Đoạn text không có dấu câu đặc biệt hoặc là newline
99
+ sep = '\n' if '\n' in parts[i] else ' '
100
+ sentences.append((parts[i].strip(), sep))
101
+ i += 1
102
+ else:
103
+ i += 1
104
+
105
+ chunks = []
106
+ current_chunk = []
107
+ word_count = 0
108
+ last_separator = ' '
109
+
110
+ for sentence, separator in sentences:
111
+ words_in_sentence = sentence.split()
112
+
113
+ if word_count + len(words_in_sentence) > max_words:
114
+ if current_chunk:
115
+ chunks.append((" ".join(current_chunk), last_separator))
116
+ current_chunk = words_in_sentence
117
+ word_count = len(words_in_sentence)
118
+ else:
119
+ current_chunk.extend(words_in_sentence)
120
+ word_count += len(words_in_sentence)
121
+
122
+ last_separator = separator
123
+
124
+ if current_chunk:
125
+ chunks.append((" ".join(current_chunk), last_separator))
126
+
127
+ return chunks
128
+
129
+ def translate_long_text(self, text: str, max_new_tokens: int = 2048) -> str:
130
+ """
131
+ Dịch văn bản dài, GHỮ NGUYÊN dấu phân cách gốc khi ghép.
132
+ """
133
+ chunks = self._split_into_chunks(text, max_words=150)
134
+ translated_parts = []
135
+
136
+ for i, (chunk, separator) in enumerate(chunks):
137
+ translated = self.translate(chunk, max_new_tokens=max_new_tokens)
138
+
139
+ # Ghép chunk với dấu phân cách phù hợp
140
+ if i < len(chunks) - 1: # Không phải chunk cuối
141
+ if separator == '\n':
142
+ translated_parts.append(translated + '\n')
143
+ elif separator in '.!?':
144
+ translated_parts.append(translated + ' ') # Dấu câu đã có sẵn trong translated
145
+ else:
146
+ translated_parts.append(translated + ' ')
147
+ else: # Chunk cuối
148
+ translated_parts.append(translated)
149
+
150
+ return ''.join(translated_parts)
readme_md.md ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: MedCrab Translation
3
+ emoji: 🦀
4
+ colorFrom: purple
5
+ colorTo: blue
6
+ sdk: gradio
7
+ sdk_version: 5.0.0
8
+ app_file: app.py
9
+ pinned: false
10
+ license: cc-by-nc-4.0
11
+ python_version: 3.10
12
+ models:
13
+ - deepseek-ai/DeepSeek-OCR
14
+ tags:
15
+ - medical
16
+ - translation
17
+ - ocr
18
+ - vietnamese
19
+ ---
20
+
21
+ # 🦀 MedCrab Translation
22
+
23
+ Ứng dụng quét OCR tài liệu y khoa và dịch trực tiếp sang tiếng Việt với hiệu ứng streaming.
24
+
25
+ ## Tính năng
26
+
27
+ - 📄 Hỗ trợ PDF và hình ảnh
28
+ - 🔍 OCR chính xác với DeepSeek-OCR
29
+ - 🦀 Dịch y khoa chuyên sâu với MedCrab
30
+ - ⚡ Streaming real-time
31
+ - 🎨 Giao diện thân thiện
32
+
33
+ ## Sử dụng
34
+
35
+ 1. Tải lên file PDF hoặc hình ảnh y khoa
36
+ 2. Chọn số trang (nếu là PDF)
37
+ 3. Chọn chế độ OCR
38
+ 4. Nhấn "Quét OCR + Dịch tiếng Việt"
39
+
40
+ ## Chế độ OCR
41
+
42
+ - **Crab**: Chế độ cân bằng (khuyên dùng)
43
+ - **Base**: Chế độ nhanh
44
+
45
+ ## Yêu cầu GPU
46
+
47
+ Space này cần GPU để chạy. Hugging Face cung cấp GPU miễn phí với giới hạn thời gian sử dụng.
48
+
49
+ ## License
50
+
51
+ Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0)
52
+
53
+ This work is licensed under CC BY-NC 4.0. You are free to:
54
+ - Share: copy and redistribute the material
55
+ - Adapt: remix, transform, and build upon the material
56
+
57
+ Under the following terms:
58
+ - Attribution: You must give appropriate credit
59
+ - NonCommercial: You may not use the material for commercial purposes
60
+
61
+ See: https://creativecommons.org/licenses/by-nc/4.0/
requirements.txt ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch==2.6.0
2
+ transformers==4.46.3
3
+ tokenizers==0.20.3
4
+ accelerate
5
+ einops
6
+ addict
7
+ easydict
8
+ torchvision
9
+ flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
10
+ PyMuPDF
11
+ hf_transfer