feat(experimental): extend IDocTags tokens (#439)

feat: extend IDocTags tokens [WIP] Signed-off-by: Panos Vagenas <pva@zurich.ibm.com>

Author: vagenas

docling_core/experimental/idoctags.py

@@ -1,5 +1,6 @@
 """Define classes for DocTags serialization."""
 
+from enum import Enum
 from typing import Any, Final, Optional
 from xml.dom.minidom import parseString
 
@@ -10,13 +11,15 @@
  BaseDocSerializer,
  BaseMetaSerializer,
  BasePictureSerializer,
+ BaseTableSerializer,
  SerializationResult,
 )
 from docling_core.transforms.serializer.common import create_ser_result
 from docling_core.transforms.serializer.doctags import (
  DocTagsDocSerializer,
  DocTagsParams,
  DocTagsPictureSerializer,
+ DocTagsTableSerializer,
  _get_delim,
  _wrap,
 )
@@ -40,6 +43,25 @@
 DOCTAGS_VERSION: Final = "1.0.0"
 
 
+class IDocTagsTableToken(str, Enum):
+ """Class to represent an LLM friendly representation of a Table."""
+
+ CELL_LABEL_COLUMN_HEADER = "<column_header/>"
+ CELL_LABEL_ROW_HEADER = "<row_header/>"
+ CELL_LABEL_SECTION_HEADER = "<shed/>"
+ CELL_LABEL_DATA = "<data/>"
+
+ OTSL_ECEL = "<ecel/>" # empty cell
+ OTSL_FCEL = "<fcel/>" # cell with content
+ OTSL_LCEL = "<lcel/>" # left looking cell,
+ OTSL_UCEL = "<ucel/>" # up looking cell,
+ OTSL_XCEL = "<xcel/>" # 2d extension cell (cross cell),
+ OTSL_NL = "<nl/>" # new line,
+ OTSL_CHED = "<ched/>" # - column header cell,
+ OTSL_RHED = "<rhed/>" # - row header cell,
+ OTSL_SROW = "<srow/>" # - section row cell
+
+
 class IDocTagsParams(DocTagsParams):
  """DocTags-specific serialization parameters."""
 
@@ -166,6 +188,7 @@ def serialize(
  temp_doc,
  add_cell_location=False,
  self_closing=params.do_self_closing,
+ table_token=IDocTagsTableToken,
  )
  body += otsl_content
  res_parts.append(create_ser_result(text=body, span_source=item))
@@ -184,11 +207,19 @@ def serialize(
  return create_ser_result(text=text_res, span_source=res_parts)
 
 
+class IDocTagsTableSerializer(DocTagsTableSerializer):
+ """DocTags-specific table item serializer."""
+
+ def _get_table_token(self) -> Any:
+ return IDocTagsTableToken
+
+
 class IDocTagsDocSerializer(DocTagsDocSerializer):
  """DocTags document serializer."""
 
  picture_serializer: BasePictureSerializer = IDocTagsPictureSerializer()
  meta_serializer: BaseMetaSerializer = IDocTagsMetaSerializer()
+ table_serializer: BaseTableSerializer = IDocTagsTableSerializer()
  params: IDocTagsParams = IDocTagsParams()
 
  @override
@@ -207,7 +238,7 @@ def serialize_doc(
  text_res = delim.join([p.text for p in parts if p.text])
 
  if self.params.add_page_break:
- page_sep = f"<{DocumentToken.PAGE_BREAK.value}>"
+ page_sep = f"<{DocumentToken.PAGE_BREAK.value}{'/' if self.params.do_self_closing else ''}>"
  for full_match, _, _ in self._get_page_breaks(text=text_res):
  text_res = text_res.replace(full_match, page_sep)
 

docling_core/transforms/serializer/doctags.py

@@ -50,7 +50,7 @@
  TextItem,
 )
 from docling_core.types.doc.labels import DocItemLabel, PictureClassificationLabel
-from docling_core.types.doc.tokens import DocumentToken
+from docling_core.types.doc.tokens import DocumentToken, TableToken
 
 
 def _wrap(text: str, wrap_tag: str) -> str:
@@ -147,6 +147,7 @@ def serialize(
  if isinstance(item, CodeItem):
  language_token = DocumentToken.get_code_language_token(
  code_language=item.code_language,
+ self_closing=params.do_self_closing,
  )
  text_part = f"{language_token}{text_part}"
  else:
@@ -171,6 +172,9 @@ def serialize(
 class DocTagsTableSerializer(BaseTableSerializer):
  """DocTags-specific table item serializer."""
 
+ def _get_table_token(self) -> Any:
+ return TableToken
+
  @override
  def serialize(
  self,
@@ -203,6 +207,7 @@ def serialize(
  xsize=params.xsize,
  ysize=params.ysize,
  visited=visited,
+ table_token=self._get_table_token(),
  )
  res_parts.append(create_ser_result(text=otsl_text, span_source=item))
 

docling_core/types/doc/document.py

@@ -2267,6 +2267,8 @@ def export_to_otsl(
 
  from docling_core.transforms.serializer.doctags import DocTagsDocSerializer
 
+ table_token = kwargs.get("table_token", TableToken)
+
  doc_serializer = DocTagsDocSerializer(doc=doc)
  body = []
  nrows = self.data.num_rows
@@ -2309,34 +2311,34 @@ def export_to_otsl(
  if rowstart == i and colstart == j:
  if len(content) > 0:
  if cell.column_header:
- body.append(str(TableToken.OTSL_CHED.value))
+ body.append(str(table_token.OTSL_CHED.value))
  elif cell.row_header:
- body.append(str(TableToken.OTSL_RHED.value))
+ body.append(str(table_token.OTSL_RHED.value))
  elif cell.row_section:
- body.append(str(TableToken.OTSL_SROW.value))
+ body.append(str(table_token.OTSL_SROW.value))
  else:
- body.append(str(TableToken.OTSL_FCEL.value))
+ body.append(str(table_token.OTSL_FCEL.value))
  if add_cell_location:
  body.append(str(cell_loc))
  if add_cell_text:
  body.append(str(content))
  else:
- body.append(str(TableToken.OTSL_ECEL.value))
+ body.append(str(table_token.OTSL_ECEL.value))
  else:
  add_cross_cell = False
  if rowstart != i:
  if colspan == 1:
- body.append(str(TableToken.OTSL_UCEL.value))
+ body.append(str(table_token.OTSL_UCEL.value))
  else:
  add_cross_cell = True
  if colstart != j:
  if rowspan == 1:
- body.append(str(TableToken.OTSL_LCEL.value))
+ body.append(str(table_token.OTSL_LCEL.value))
  else:
  add_cross_cell = True
  if add_cross_cell:
- body.append(str(TableToken.OTSL_XCEL.value))
- body.append(str(TableToken.OTSL_NL.value))
+ body.append(str(table_token.OTSL_XCEL.value))
+ body.append(str(table_token.OTSL_NL.value))
  body_str = "".join(body)
  return body_str
 

docling_core/types/doc/tokens.py

@@ -257,9 +257,12 @@ def get_picture_classification_token(classification: str) -> str:
  return _PictureClassificationToken(f"<{classification}>").value
 
  @staticmethod
- def get_code_language_token(code_language: str) -> str:
+ def get_code_language_token(code_language: str, self_closing: bool = False) -> str:
  """Function to get the token for a given code language."""
- return _CodeLanguageToken(f"<_{code_language}_>").value
+ if self_closing:
+ return f"<{code_language}/>"
+ else:
+ return _CodeLanguageToken(f"<_{code_language}_>").value
 
  @staticmethod
  def get_location_token(