doc-page-extractor 0.0.10__py3-none-any.whl → 0.1.0__py3-none-any.whl

doc_page_extractor/__init__.py

@@ -1,5 +1,15 @@
 from .extractor import DocExtractor
 from .clipper import clip, clip_from_image
 from .plot import plot
-from .types import ExtractedResult, OCRFragment, LayoutClass, Layout
-from .rectangle import Point, Rectangle
+from .rectangle import Point, Rectangle
+from .types import (
+  ExtractedResult,
+  OCRFragment,
+  LayoutClass,
+  TableLayoutParsedFormat,
+  Layout,
+  BaseLayout,
+  PlainLayout,
+  FormulaLayout,
+  TableLayout,
+)

doc_page_extractor/extractor.py

@@ -1,6 +1,6 @@
 import os
 
-from typing import Literal
+from typing import Literal, Generator
 from pathlib import Path
 from PIL.Image import Image
 from doclayout_yolo import YOLOv10
@@ -9,10 +9,22 @@ from .ocr import OCR
 from .ocr_corrector import correct_fragments
 from .raw_optimizer import RawOptimizer
 from .rectangle import intersection_area, Rectangle
-from .types import ExtractedResult, OCRFragment, LayoutClass, Layout
 from .downloader import download
+from .table import Table
+from .latex import LaTeX
 from .layout_order import LayoutOrder
 from .overlap import merge_fragments_as_line, remove_overlap_layouts
+from .clipper import clip_from_image
+from .types import (
+  ExtractedResult,
+  OCRFragment,
+  TableLayoutParsedFormat,
+  Layout,
+  LayoutClass,
+  PlainLayout,
+  TableLayout,
+  FormulaLayout,
+)
 
 
 class DocExtractor:
@@ -21,12 +33,26 @@ class DocExtractor:
       model_dir_path: str,
       device: Literal["cpu", "cuda"] = "cpu",
       ocr_for_each_layouts: bool = True,
+      extract_formula: bool = True,
+      extract_table_format: TableLayoutParsedFormat | None = None,
     ):
     self._model_dir_path: str = model_dir_path
     self._device: Literal["cpu", "cuda"] = device
     self._ocr_for_each_layouts: bool = ocr_for_each_layouts
-    self._ocr: OCR = OCR(device, model_dir_path)
+    self._extract_formula: bool = extract_formula
+    self._extract_table_format: TableLayoutParsedFormat | None = extract_table_format
     self._yolo: YOLOv10 | None = None
+    self._ocr: OCR = OCR(
+      device=device,
+      model_dir_path=os.path.join(model_dir_path, "onnx_ocr"),
+    )
+    self._table: Table = Table(
+      device=device,
+      model_path=os.path.join(model_dir_path, "struct_eqtable"),
+    )
+    self._latex: LaTeX = LaTeX(
+      model_path=os.path.join(model_dir_path, "latex"),
+    )
     self._layout_order: LayoutOrder = LayoutOrder(
       model_path=os.path.join(model_dir_path, "layoutreader"),
     )
@@ -40,7 +66,7 @@ class DocExtractor:
     raw_optimizer = RawOptimizer(image, adjust_points)
     fragments = list(self._ocr.search_fragments(raw_optimizer.image_np))
     raw_optimizer.receive_raw_fragments(fragments)
-    layouts = self._get_layouts(raw_optimizer.image)
+    layouts = list(self._yolo_extract_layouts(raw_optimizer.image))
     layouts = self._layouts_matched_by_fragments(fragments, layouts)
     layouts = remove_overlap_layouts(layouts)
 
@@ -50,6 +76,8 @@ class DocExtractor:
     layouts = self._layout_order.sort(layouts, raw_optimizer.image.size)
     layouts = [layout for layout in layouts if self._should_keep_layout(layout)]
 
+    self._parse_table_and_formula_layouts(layouts, raw_optimizer)
+
     for layout in layouts:
       layout.fragments = merge_fragments_as_line(layout.fragments)
 
@@ -62,7 +90,7 @@ class DocExtractor:
       adjusted_image=raw_optimizer.adjusted_image,
     )
 
-  def _get_layouts(self, source: Image) -> list[Layout]:
+  def _yolo_extract_layouts(self, source: Image) -> Generator[Layout, None, None]:
     # about source parameter to see:
     # https://github.com/opendatalab/DocLayout-YOLO/blob/7c4be36bc61f11b67cf4a44ee47f3c41e9800a91/doclayout_yolo/data/build.py#L157-L175
     det_res = self._get_yolo().predict(
@@ -72,7 +100,6 @@ class DocExtractor:
       device=self._device    # Device to use (e.g., "cuda" or "cpu")
     )
     boxes = det_res[0].__dict__["boxes"]
-    layouts: list[Layout] = []
 
     for cls_id, rect in zip(boxes.cls, boxes.xyxy):
       cls_id = cls_id.item()
@@ -89,9 +116,12 @@ class DocExtractor:
         lb=(x1, y2),
         rb=(x2, y2),
       )
-      layouts.append(Layout(cls, rect, []))
-
-    return layouts
+      if cls == LayoutClass.TABLE:
+        yield TableLayout(cls=cls, rect=rect, fragments=[], parsed=None)
+      elif cls == LayoutClass.ISOLATE_FORMULA:
+        yield FormulaLayout(cls=cls, rect=rect, fragments=[], latex=None)
+      else:
+        yield PlainLayout(cls=cls, rect=rect, fragments=[])
 
   def _layouts_matched_by_fragments(self, fragments: list[OCRFragment], layouts: list[Layout]):
     layouts_group = self._split_layouts_by_group(layouts)
@@ -107,6 +137,17 @@ class DocExtractor:
     for layout in layouts:
       correct_fragments(self._ocr, source, layout)
 
+  def _parse_table_and_formula_layouts(self, layouts: list[Layout], raw_optimizer: RawOptimizer):
+    for layout in layouts:
+      if isinstance(layout, FormulaLayout) and self._extract_formula:
+        image = clip_from_image(raw_optimizer.image, layout.rect)
+        layout.latex = self._latex.extract(image)
+      elif isinstance(layout, TableLayout) and self._extract_table_format is not None:
+        image = clip_from_image(raw_optimizer.image, layout.rect)
+        parsed = self._table.predict(image, self._extract_table_format)
+        if parsed is not None:
+          layout.parsed = (parsed, self._extract_table_format)
+
   def _split_layouts_by_group(self, layouts: list[Layout]):
     texts_layouts: list[Layout] = []
     abandon_layouts: list[Layout] = []
@@ -149,9 +190,11 @@ class DocExtractor:
 
   def _get_yolo(self) -> YOLOv10:
     if self._yolo is None:
+      base_path = os.path.join(self._model_dir_path, "yolo")
+      os.makedirs(base_path, exist_ok=True)
       yolo_model_url = "https://huggingface.co/opendatalab/PDF-Extract-Kit-1.0/resolve/main/models/Layout/YOLO/doclayout_yolo_ft.pt"
       yolo_model_name = "doclayout_yolo_ft.pt"
-      yolo_model_path = Path(os.path.join(self._model_dir_path, yolo_model_name))
+      yolo_model_path = Path(os.path.join(base_path, yolo_model_name))
       if not yolo_model_path.exists():
         download(yolo_model_url, yolo_model_path)
       self._yolo = YOLOv10(str(yolo_model_path))

doc_page_extractor/latex.py

@@ -0,0 +1,57 @@
+import os
+import torch
+import requests
+
+from munch import Munch
+from pix2tex.cli import LatexOCR
+from PIL.Image import Image
+from .utils import expand_image
+
+
+class LaTeX:
+  def __init__(self, model_path: str):
+    self._model_path: str = model_path
+    self._model: LatexOCR | None = None
+
+  def extract(self, image: Image) -> str | None:
+    image = expand_image(image, 0.1) # 添加边缘提高识别准确率
+    model = self._get_model()
+    with torch.no_grad():
+      return model(image)
+
+  def _get_model(self) -> LatexOCR:
+    if self._model is None:
+      if not os.path.exists(self._model_path):
+        self._download_model()
+
+      self._model = LatexOCR(Munch({
+        "config": os.path.join("settings", "config.yaml"),
+        "checkpoint": os.path.join(self._model_path, "weights.pth"),
+        "no_cuda": True,
+        "no_resize": False,
+      }))
+    return self._model
+
+  # from https://github.com/lukas-blecher/LaTeX-OCR/blob/5c1ac929bd19a7ecf86d5fb8d94771c8969fcb80/pix2tex/model/checkpoints/get_latest_checkpoint.py#L37-L45
+  def _download_model(self):
+    os.makedirs(self._model_path, exist_ok=True)
+    tag = "v0.0.1"
+    files: list[tuple[str, str]] = (
+      ("weights.pth", f"https://github.com/lukas-blecher/LaTeX-OCR/releases/download/{tag}/weights.pth"),
+      ("image_resizer.pth", f"https://github.com/lukas-blecher/LaTeX-OCR/releases/download/{tag}/image_resizer.pth")
+    )
+    for file_name, url in files:
+      file_path = os.path.join(self._model_path, file_name)
+      try:
+        with open(file_path, "wb") as file:
+          response = requests.get(url, stream=True, timeout=15)
+          response.raise_for_status()
+          for chunk in response.iter_content(chunk_size=8192):
+            if chunk:  # 过滤掉保持连接的新块
+              file.write(chunk)
+          file.flush()
+
+      except BaseException as e:
+        if os.path.exists(file_path):
+          os.remove(file_path)
+        raise e

doc_page_extractor/struct_eqtable/__init__.py

@@ -0,0 +1,49 @@
+from .pix2s import Pix2Struct, Pix2StructTensorRT
+from .internvl import InternVL, InternVL_LMDeploy
+
+from transformers import AutoConfig
+
+
+__ALL_MODELS__ = {
+    'Pix2Struct': Pix2Struct,
+    'Pix2StructTensorRT': Pix2StructTensorRT,
+    'InternVL': InternVL,
+    'InternVL_LMDeploy': InternVL_LMDeploy,
+}
+
+
+def get_model_name(model_path):
+    model_config = AutoConfig.from_pretrained(
+        model_path,
+        trust_remote_code=True,
+    )
+
+    if 'Pix2Struct' in model_config.architectures[0]:
+        model_name = 'Pix2Struct'
+    elif 'InternVL' in model_config.architectures[0]:
+        model_name = 'InternVL'
+    else:
+        raise ValueError(f"Unsupported model type: {model_config.architectures[0]}")
+
+    return model_name
+
+
+def build_model(
+        model_ckpt='U4R/StructTable-InternVL2-1B',
+        cache_dir=None,
+        local_files_only=None,
+        **kwargs,
+    ):
+    model_name = get_model_name(model_ckpt)
+    if model_name == 'InternVL' and kwargs.get('lmdeploy', False):
+        model_name = 'InternVL_LMDeploy'
+    elif model_name == 'Pix2Struct' and kwargs.get('tensorrt_path', None):
+        model_name = 'Pix2StructTensorRT'
+
+    model = __ALL_MODELS__[model_name](
+        model_ckpt,
+        cache_dir=cache_dir,
+        local_files_only=local_files_only,
+        **kwargs
+    )
+    return model

doc_page_extractor/struct_eqtable/internvl/__init__.py

@@ -0,0 +1,2 @@
+from .internvl import InternVL
+from .internvl_lmdeploy import InternVL_LMDeploy

doc_page_extractor/struct_eqtable/internvl/conversation.py

@@ -0,0 +1,394 @@
+"""
+Conversation prompt templates.
+
+We kindly request that you import fastchat instead of copying this file if you wish to use it.
+If you have changes in mind, please contribute back so the community can benefit collectively and continue to maintain these valuable templates.
+"""
+
+import dataclasses
+from enum import IntEnum, auto
+from typing import Any, Dict, List, Tuple, Union
+
+
+class SeparatorStyle(IntEnum):
+    """Separator styles."""
+
+    ADD_COLON_SINGLE = auto()
+    ADD_COLON_TWO = auto()
+    ADD_COLON_SPACE_SINGLE = auto()
+    NO_COLON_SINGLE = auto()
+    NO_COLON_TWO = auto()
+    ADD_NEW_LINE_SINGLE = auto()
+    LLAMA2 = auto()
+    CHATGLM = auto()
+    CHATML = auto()
+    CHATINTERN = auto()
+    DOLLY = auto()
+    RWKV = auto()
+    PHOENIX = auto()
+    ROBIN = auto()
+    FALCON_CHAT = auto()
+    CHATGLM3 = auto()
+    INTERNVL_ZH = auto()
+    MPT = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that manages prompt templates and keeps all conversation history."""
+
+    # The name of this template
+    name: str
+    # The template of the system prompt
+    system_template: str = '{system_message}'
+    # The system message
+    system_message: str = ''
+    # The names of two roles
+    roles: Tuple[str] = ('USER', 'ASSISTANT')
+    # All messages. Each item is (role, message).
+    messages: List[List[str]] = ()
+    # The number of few shot examples
+    offset: int = 0
+    # The separator style and configurations
+    sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
+    sep: str = '\n'
+    sep2: str = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: Union[str, List[str]] = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: List[int] = None
+
+    def get_prompt(self) -> str:
+        """Get the prompt for generation."""
+        system_prompt = self.system_template.format(system_message=self.system_message)
+        if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ': '  # must be end with a space
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + message + seps[i % 2]
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.RWKV:
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += (
+                        role
+                        + ': '
+                        + message.replace('\r\n', '\n').replace('\n\n', '\n')
+                    )
+                    ret += '\n\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA2:
+            seps = [self.sep, self.sep2]
+            if self.system_message:
+                ret = system_prompt
+            else:
+                ret = '[INST] '
+            for i, (role, message) in enumerate(self.messages):
+                tag = self.roles[i % 2]
+                if message:
+                    if i == 0:
+                        ret += message + ' '
+                    else:
+                        ret += tag + ' ' + message + seps[i % 2]
+                else:
+                    ret += tag
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM:
+            # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
+            # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
+            round_add_n = 1 if self.name == 'chatglm2' else 0
+            if system_prompt:
+                ret = system_prompt + self.sep
+            else:
+                ret = ''
+
+            for i, (role, message) in enumerate(self.messages):
+                if i % 2 == 0:
+                    ret += f'[Round {i//2 + round_add_n}]{self.sep}'
+
+                if message:
+                    ret += f'{role}：{message}{self.sep}'
+                else:
+                    ret += f'{role}：'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATML:
+            ret = '' if system_prompt == '' else system_prompt + self.sep + '\n'
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep + '\n'
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM3:
+            ret = ''
+            if self.system_message:
+                ret += system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + ' ' + message
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATINTERN:
+            # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                # if i % 2 == 0:
+                #     ret += "<s>"
+                if message:
+                    ret += role + ':' + message + seps[i % 2] + '\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.DOLLY:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ':\n' + message + seps[i % 2]
+                    if i % 2 == 1:
+                        ret += '\n\n'
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.PHOENIX:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + '<s>' + message + '</s>'
+                else:
+                    ret += role + ': ' + '<s>'
+            return ret
+        elif self.sep_style == SeparatorStyle.ROBIN:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ':\n' + message + self.sep
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.FALCON_CHAT:
+            ret = ''
+            if self.system_message:
+                ret += system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+
+            return ret
+        elif self.sep_style == SeparatorStyle.INTERNVL_ZH:
+            seps = [self.sep, self.sep2]
+            ret = self.system_message + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        else:
+            raise ValueError(f'Invalid style: {self.sep_style}')
+
+    def set_system_message(self, system_message: str):
+        """Set the system message."""
+        self.system_message = system_message
+
+    def append_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def update_last_message(self, message: str):
+        """Update the last output.
+
+        The last message is typically set to be None when constructing the prompt,
+        so we need to update it in-place after getting the response from a model.
+        """
+        self.messages[-1][1] = message
+
+    def to_gradio_chatbot(self):
+        """Convert the conversation to gradio chatbot format."""
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        """Convert the conversation to OpenAI chat completion format."""
+        ret = [{'role': 'system', 'content': self.system_message}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append({'role': 'user', 'content': msg})
+            else:
+                if msg is not None:
+                    ret.append({'role': 'assistant', 'content': msg})
+        return ret
+
+    def copy(self):
+        return Conversation(
+            name=self.name,
+            system_template=self.system_template,
+            system_message=self.system_message,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            stop_str=self.stop_str,
+            stop_token_ids=self.stop_token_ids,
+        )
+
+    def dict(self):
+        return {
+            'template_name': self.name,
+            'system_message': self.system_message,
+            'roles': self.roles,
+            'messages': self.messages,
+            'offset': self.offset,
+        }
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert (
+            template.name not in conv_templates
+        ), f'{template.name} has been registered.'
+
+    conv_templates[template.name] = template
+
+
+def get_conv_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name].copy()
+
+
+# Both Hermes-2 and internlm2-chat are chatml-format conversation templates. The difference
+# is that during training, the preprocessing function for the Hermes-2 template doesn't add
+# <s> at the beginning of the tokenized sequence, while the internlm2-chat template does.
+# Therefore, they are completely equivalent during inference.
+register_conv_template(
+    Conversation(
+        name='Hermes-2',
+        system_template='<|im_start|>system\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        # system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        system_message='You are a Table Image to LaTeX/Markdown/HMTL Code converter.',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+        stop_token_ids=[
+            2,
+            6,
+            7,
+            8,
+        ],
+        stop_str='<|endoftext|>',
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name='internlm2-chat',
+        system_template='<|im_start|>system\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+        stop_token_ids=[
+            2,
+            92543,
+            92542
+        ]
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name='phi3-chat',
+        system_template='<|system|>\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|user|>\n', '<|assistant|>\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|end|>',
+        stop_token_ids=[
+            2,
+            32000,
+            32007
+        ]
+    )
+)