
    4 jJS                         d dl Z d dlZd dlmZ d dlmc mZ ej                  Z	 ej                  de	      Zd Z G d d      Zd Zd Zd Zd	 Zd
 Zd Zd Zd Zy)    N)groupbyconsolidatec                   0    e Zd ZdZdZdZdZdZdZdZ	dZ
d	Zy
)bcolorsz[95mz[94mz[96mz[92mz[93mz[91mz[0mz[1mz[4mN)__name__
__module____qualname__HEADEROKBLUEOKCYANOKGREENWARNINGFAILENDCBOLD	UNDERLINE     V/home/agent/.local/lib/python3.12/site-packages/docling_ibm_models/tableformer/otsl.pyr   r      s/    FFFGGDDDIr   r   c                 L    g }g d}| D ]  }||vs|j                  |        |S )N)z<pad>z<unk>z<start>z<end>)append)rs_listnew_rs_list	stop_listtags       r   
otsl_cleanr      s9    K6I $is#$ r   c                    t        | d       D cg c]  \  }}|r	t        |       }}}d}t        |      dkD  rt        |d         dz   }t        |      D ]'  \  }}	|	j	                  d       t        |	      |k7  s&d}) |rE|rAt
        j                  dj                  t        j                  t        j                               |S d|z   d	z   }
t
        j                  |
j                  t        j                  t        j                               t
        j                  d
j                  t        j                  t        j                               |S c c}}w )Nc                     | dk(  S Nnlr   xs    r   <lambda>zotsl_sqr_chk.<locals>.<lambda>%   
    qDy r   Tr      r    Fz{}*OK* Table is square! *OK*{}z{}*ERR* z *ERR*{}z${}*ERR* Table is not square! *ERR*{})r   listlen	enumerater   loggerdebugformatr   r   r   r   )r   namelogdebugkgrouprs_list_splitisSquareinit_tag_lenindlnerr_names              r   otsl_sqr_chkr6   #   s$   $+G5H$I 5QRUM  H
=A=+,q0 / 	!GCIIdO2w,& 	! 4;; O "D(:5HLLw||DELL6==LL',,
 O3s
   
EEc                 ^   g }t        | d       D cg c]  \  }}|r	t        |       }}}d}t        |      D ]  \  }}t        |      |kD  st        |      }! t        |      D ]<  \  }}||g|t        |      z
  z  z  }|j	                  d       |j                  |       > |S c c}}w )Nc                     | dk(  S r   r   r!   s    r   r#   z!otsl_pad_to_sqr.<locals>.<lambda>C   r$   r   r   r    )r   r&   r(   r'   r   extend)	r   pad_tagnew_listr.   r/   r0   max_row_lenr3   r4   s	            r   otsl_pad_to_sqrr=   @   s    H$+G5H$I 5QRUM  K]+ "Rr7[ b'K" ]+ R
wi;R011
		$ Os
   
B)B)c                 h   d}d}| D ]  }|dv s|dz  } |t        |      k7  rd|z   dz   }t        j                  |j                  t        j
                  t        j                               d}t        j                  |j                  t        j
                  t        j                               d}|S )	Nr   Tfcelchedrhedsrowecelr%   z	{}*!ERR* z	 *ERR!*{}z2{}*!ERR* Tags are not in sync with cells! *ERR!*{}F)r'   r)   r*   r+   r   r   r   )	r   cellsr,   r-   countCellTagsisGoodrsTagr5   err_msgs	            r   otsl_tags_cells_sync_chkrJ   P   s    MF <<QM E
"%3X__W\\7<<@AFW^^GLL',,?@Mr   c                     d}d}g d}||vr9|t        |       dz
  k  r(|dz  }|dz  }| |   |   }||vr|t        |       dz
  k  r(||v r|dz  }|S )Nr%   ucel)r@   rA   rB   rC   rD   lcelr    r'   rs_splitr"   ydistanceelemgoodlists         r   otsl_check_downrU   _   sz    HDEH
h
1s8}q'8#8	QA{1~ h
1s8}q'8#8 xAOr   c                     d}d}g d}||vr?|t        | |         dz
  k  r+|dz  }|dz  }| |   |   }||vr|t        | |         dz
  k  r+||v r|dz  }|S )Nr%   rM   )r@   rA   rB   rC   rD   rL   r    rN   rO   s         r   otsl_check_rightrW   l   s    HDEH
h
1HQK(81(<#=	QA{1~ h
1HQK(81(<#= xAOr   c                 (   t        |       dk(  rg S | d   dvr| S g }|rAt        j                  dj                  t        j
                  t        j                               t        | d|      sMt        j                  dj                  t        j
                  t        j                               t        | d      } t        | d       D cg c]  \  }}|r	t        |       }}}|rt        j                  d       g }d	}t        |      D ]  \  }}	g }
|sd
||   v r|
j                  d       d}|rd
||   vr|
j                  d       d	}|
j                  d       t        ||         D ]8  \  }}|dv sd}d}d}d}d	}|dz   t        ||         k  r||   |dz      dk(  rt        |||      }d}|dz   t        |      k  r||dz      |   dk(  rt        |||      }d}|dz   t        ||         k  r||   |dz      dk(  rt        |||      }t        |||      }d}t        |||z         D ]7  }t        |||z         D ]#  }t!        |      dz   t!        |      z   }||v s"d	}% 9 |rSt        |||z         D ]A  }t        |||z         D ]-  }t!        |      dz   t!        |      z   }|j                  |       / C |r|
j                  d       |dkD  r |
j                  dt!        |      z   dz          |dkD  r |
j                  dt!        |      z   dz          |dkD  r@|
j                  dt!        |      z   dz          |
j                  dt!        |      z   dz          |
j                  d       |
j                  d       |
j                  d       |
j                  d       ; |
j                  d       |j#                  |
        |r?t        j                  d       t        j                  |       t        j                  d       |S c c}}w )Nr   r?   z{}*Reconstructing HTML...*{}z---z{}*Padding to square...*{}rM   c                     | dk(  S r   r   r!   s    r   r#   zotsl_to_html.<locals>.<lambda>   r$   r    FrA   <thead>Tz</thead>z<tr>r%   rL   xcel_<tdz
 colspan=""z
 rowspan=">z</td><td></tr>zD*********************** registry_2d_span ***************************zD********************************************************************)r'   r)   r*   r+   r   r   r   r6   r=   r   r&   r(   r   rW   rU   rangestrr9   )r   r-   
html_tabler.   r/   r0   registry_2d_spanthead_present
rs_row_indrs_row	html_listrs_cell_indrs_cellrdistddistxrdistxddistspanr"   rQ   reg2dinds                        r   otsl_to_htmlrs   y   sP   
7|q	qzAAJ*11'//7<<P	
 1188',,WX!'62 %,G5H$I 5QRUM  R 
 M'6 K%
F	z22  + $]:66  , % $-mJ.G$H ;	. KBB?Sz)B%CC$Z0qAVK 0Z X#>C$66 %Z!^4[AVK /{J W#?Sz)B%CC$Z0qAVK!1);
" "1Z!X#
 "'{F[4H!I 1A%*:v
7J%K 1+.q6C<#a&+@#+/?#?+0D11   &+;8L%M F).z6J;N)O !FA/21v|c!f/DH$4$;$;H$E!FF $$U+qy!((E
)BS)HIqy!((E
)BS)HIz!((F)Cc)IJ!((F)Cc)IJ$$S)$$W-$$V,$$W-w;	.x 	!)$WK%Z R	
 	%&R	
 Ks   	
PPc                 X   t        j                  | d   d         }| }|r||d   d<   dg i|d   d<   dg i|d   d<   i dddd	d
dddddddddddddddddddddd d!d"d#d$d%d&d'd(d)d*}i d+dd,d	d-dd.dd/dd0dd1dd2dd3dd4dd5dd6dd7dd8d!d9d#d:d%d;d'd(d)d<}	g }
g }g }i }d=}d=}d=}g }g }|r@t        j                  d>       t        j                  d?j	                  | d@   | dA                d=}|rBt        j                  dB       t        j                  |d          t        j                  dC       |r?t        j                  dD       t        j                  dE       t        j                  dD       t        t        |d               D ]  }|d   |   }d>}d>}|d=kD  r|d   |dFz
     }|t        |d         dFz
  k  r|d   |dFz      }|dGvsD|dHv r|d=kD  r|t        ||dFz
           k\  rdIi fc S ||dFz
     |   }|dF   d=kD  r|d=   d=k(  r0|j                  dJ       |j                  d=|dF   dFz
  g       |dFz  }nHt        |d=         D ]7  }|j                  dK       |j                  |d=   |z
  |dF   dFz
  g       |dFz  }9 ||dFz
     |   }|dF   d=kD  r|dLv r|d=kD  rt        |      }t        ||dFz
           }||k  r||z
  dFz
  }|d=kD  r|r$t        j                  dMj	                  |             t        |      D ]p  }||dFz
     ||z      }|d=   d=k(  r+|j                  dJ       |j                  d=|dF   dFz
  g       D|j                  dK       |j                  |d=   |dF   dFz
  g       r |dNv rSdO}dP| d   v r| d   dP   |   d   }ndQ}t        t        dR |            }t        t        dS |            }t        t        dT |            }t        t        dU |            }t        t        dV |            }t        t        dW |            }t        t        dX |            }t        t        dY |            }t        t        dZ |            }t        |      d=kD  rdI}|r3g |d   dP   |   d<   |j                  d[       |j                  d=d=g       n$|j                  d\       |j                  d=d=g       |dFz  }|dFz  }|d]k(  rt        |      |kD  rt        |      }|j                  d^       |
j                  t        j                  |             |j                  t        j                  |             |rt        |       g }|j                  d_d_g       |j                  t        j                  |             g }d=}|dFz  }||v rl||	vrh||	vrd||   } | dFg||dFz
  <   | d=g|t        |      dFz
  <   t        | dFz
        D ]1  }|j                  d`       |j                  | |z
  dFz
  d=g       |dFz  }3 ||	v r-||vr)||vr%|	|   }!dF|!g||dFz
  <   d=|!dFz
  g|t        |      dFz
  <   ||	v s|	|   }!||v rg||   } | |!g||dFz
  <   | |!dFz
  g}"|"|t        |      dFz
  <   t        | dFz
        D ]/  }|j                  dK       |j                  | |z
  dFz
  |!dFz
  g       1 ||v s$||   } | |!g||dFz
  <   | |!dFz
  g}"|"|t        |      dFz
  <   t        | dFz
        D ]/  }|j                  dK       |j                  | |z
  dFz
  |!dFz
  g       1  d?j	                  | d@   | dA         }#t        ||#|      }$|$st        |d\      }dP|d   v rt        ||d   dP   |#|      }%g }&|$rt        ||      }&|&|d   d   d<   ||d   d   d<   ||da<   ||db<   t        |
      |dc<   t        |d         |dd<   t        |      |de<   |r|$r%r|j!                  |       |rt        j                  dfj	                  t"        j$                  t"        j&                               t        j                  |&       |d   d   }'t        j                  dgj	                  t"        j(                  t"        j&                               t        j                  |'       dO|fS )hz
    Converts table structure from HTML to RS

    Parameters
    ----------
    table : json
        line from jsonl
    writer : writer
        Writes lines into output jsonl
    html	structurehtml_structuretokenshtml_restored_structurez colspan="20"   z colspan="19"   z colspan="18"   z colspan="17"   z colspan="16"   z colspan="15"   z colspan="14"   z colspan="13"   z colspan="12"   z colspan="11"   z colspan="10"
   z colspan="2"   z colspan="3"   z colspan="4"   z colspan="5"   z colspan="6"   z colspan="7"      	   )z colspan="8"z colspan="9"z rowspan="20"z rowspan="19"z rowspan="18"z rowspan="17"z rowspan="16"z rowspan="15"z rowspan="14"z rowspan="13"z rowspan="12"z rowspan="11"z rowspan="10"z rowspan="2"z rowspan="3"z rowspan="4"z rowspan="5"z rowspan="6"z rowspan="7")z rowspan="8"z rowspan="9"r   rZ   z*** {}: {} ***splitfilenamezB========================== Input HTML ============================zB==================================================================z********z* OTSL *r%   )r[   z<tbody>)ra   r^   rb   FrL   r\   )rb   zExtra columns needed in row: {})ra   r^   TrE   fc                     | dk7  S )Nz<i>r   as    r   r#   zhtml_to_otsl.<locals>.<lambda>  
    AJ r   c                     | dk7  S )Nz<I>r   r   s    r   r#   zhtml_to_otsl.<locals>.<lambda>  r   r   c                     | dk7  S )Nz<b>r   r   s    r   r#   zhtml_to_otsl.<locals>.<lambda>  r   r   c                     | dk7  S )Nz<B>r   r   s    r   r#   zhtml_to_otsl.<locals>.<lambda>  r   r   c                     | dk7  S )N r   r   s    r   r#   zhtml_to_otsl.<locals>.<lambda>  s
    AH r   c                     | dk7  S )Nz</b>r   r   s    r   r#   zhtml_to_otsl.<locals>.<lambda>  
    AK r   c                     | dk7  S )Nz</B>r   r   s    r   r#   zhtml_to_otsl.<locals>.<lambda>  r   r   c                     | dk7  S )Nz</i>r   r   s    r   r#   zhtml_to_otsl.<locals>.<lambda>  r   r   c                     | dk7  S )Nz</I>r   r   s    r   r#   zhtml_to_otsl.<locals>.<lambda>  r   r   rD   r@   rb   r    rM   
otsl_spanscolsrowshtml_lenrs_lenz{}Reconstructed HTML:{}z{}Original HTML:{})copydeepcopyr)   r*   r+   rc   r'   r   r&   filterr9   printr6   r=   rJ   rs   writer   r   r   r   )(tablewriterr-   extra_debuginclude_html
use_writertable_html_structureout_linepos_colspanspos_rowspanst_cellstl_cells	t_expandstl_spanscurrent_linecurrent_columncurrent_html_cell_indcurrent_line_tagscurrent_line_expandscolnumihtml_tagprev_html_tagnext_html_tag	up_expandcicur_line_lenpre_line_lenextra_columnsclmcell_is_emptycell_tokenscolspan_lenrowspan_lennewexpt_namer1   rG   rHTMLoHTMLs(                                           r   html_to_otslr      sU     ==v{)CDH-A)*7?n23%-rNHV[! 	 		
 	 	 	 	 	 	 	 	 	 	 	  	!" 	#$ 'L, 	 		
 	 	 	 	 	 	 	 	 	 	 	  	!" 	#$ 'L, GHIHLNR%,,U7^U:=NOPFP	
 	)(34P	
 Z Z Z 3+H567 k'1!4q50:1q5AMs'12Q660:1q5AM11 33!#%Y|a7G-H)II$by( ),*: ;N KI#A,*$Q<1,-44V<077IaL1<L8MN*a/N ',IaL&9 4 1 8 8 @ 4 ; ;%.q\B%6	!q8H$I!" !/! 34 %.lQ.>$?$O	 $A,*" 9$!##&';#<L#&y1A'B#CL#l2(4|(Ca(G(1,* &$E$L$L(5%&!" (-]'; &,5lQ6F,G$03$6-"	 $-Q<1#4$5$<$<V$D$8$?$?IaLSTDT@U$V %6$<$<V$D$8$?$?)21y|a7G(H%&&" ?* $eFm+"'-"89N"O #K #&K #6*>#LM"6*>#LM"6*>#LM"6*>#LM"6*<k#JK"6*?#MN"6*?#MN"6*?#MN"6*?#MN {#a'$)M QSHV$W-.CDXN%,,V4(//A7%,,V4(//A7%*%!# 7"()F2 !23F!((.t}}->?@.? @A+,$&! %++RH5  /C!DE')$!"!<' 4$L8&28&<?JA>N!6!:;'O,S1E-F-JK #(a"8 0B-44V<077r9IA9Mq8QR*a/N0 <' 4$L8&28&<?@+>N!6!:;'!OO,S1E-F-JK <'*84 L0".}"=K;F:TH2Q67);?;FJP(-A)BQ)FG#K!O4 )008,33(2-1;?C
 !L0".}"=K;F:TH2Q67);?;FJP(-A)BQ)FG#K!O4 )008,33(2-1;?COkZ $$U7^U:5FGFHfh7H"8V4(6"")hv&w/
 EXx0<AHV./9.6HV[!(+%H\HV7|HV3H=>HZXHXX&.55goow||TUU !12)00NOU>r   )r   logging	itertoolsr   'docling_ibm_models.tableformer.settingstableformersettingssINFO	LOG_LEVELget_custom_loggerr)   
total_picsr   r   r6   r=   rJ   rU   rW   rs   r   r   r   r   <module>r      sk       3 3LL				]I	6
	 	: 

xvrr   