mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-10-30 19:33:43 +08:00
myhloli
377b09cf8c
- Remove import and usage of StructTableModel- Add support for TableMaster model- Update table model initialization logic to support TableMaster - Log error and exit if StructEqTable is selected, as it's under upgrade - Update README files to reflect changes in table parsing capabilities
29 lines
1.0 KiB
Python
29 lines
1.0 KiB
Python
from loguru import logger
|
|
|
|
try:
|
|
from struct_eqtable.model import StructTable
|
|
except ImportError:
|
|
logger.error("StructEqTable is under upgrade, the current version does not support it.")
|
|
from pypandoc import convert_text
|
|
|
|
|
|
class StructTableModel:
|
|
def __init__(self, model_path, max_new_tokens=2048, max_time=400, device = 'cpu'):
|
|
# init
|
|
self.model_path = model_path
|
|
self.max_new_tokens = max_new_tokens # maximum output tokens length
|
|
self.max_time = max_time # timeout for processing in seconds
|
|
if device == 'cuda':
|
|
self.model = StructTable(self.model_path, self.max_new_tokens, self.max_time).cuda()
|
|
else:
|
|
self.model = StructTable(self.model_path, self.max_new_tokens, self.max_time)
|
|
|
|
def image2latex(self, image) -> str:
|
|
table_latex = self.model.forward(image)
|
|
return table_latex
|
|
|
|
def image2html(self, image) -> str:
|
|
table_latex = self.image2latex(image)
|
|
table_html = convert_text(table_latex, 'html', format='latex')
|
|
return table_html
|