| 12345678910111213141516171819202122232425262728 |
- 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
|
