
    I,ic                        d Z ddlZddlZddlZddlZddlZddlmZ ddl	mc m
Z ddlmZmZ ddlmZmZmZ ddlmZmZmZ ddlmZ ddlmZ  G d d	      Z e       Zd
 Zd ZddZddZ  G d de      Z!d Z" G d dejF                        Z$ G d dejF                        Z% G d dejF                        Z&ddZ'd Z(d Z)e*dk(  r e)        yy)z
Word-MLP Sequence Model: Conv2D + LSTM for Language Prediction
ODE-CCT framework: word sequences as trajectories, next word as trajectory prediction
    N)Dataset
DataLoader)pad_sequencepack_padded_sequencepad_packed_sequence)Image	ImageDraw	ImageFont)defaultdict)Pathc                       e Zd ZdZdZdZdZdZdZdZ	dZ
dZd	Zd
ZdZ ej                   ej"                  j%                         rd      Zyd      Zy)Configz./books            N   @   MbP?g333333?   cudacpu)__name__
__module____qualname__book_folderimg_size	font_size	embed_dim
hidden_dimnum_classesseq_len
batch_sizeepochslrentropy_thresholdperiodicity_windowtorchdevicer   is_available     ex03.pyr   r      sp    K HI IJK G JF	B U\\EJJ$;$;$=&IF5IFr,   r   c                 &   g }t        |       }|j                         st        d|  d       t               S |j	                  d      D ]  }t        d|j
                          t        |ddd      5 }|j                         }d	d	d	       t        j                  d
      }|j                  |D cg c]  }|j                          c}        t        |      dk(  r
t               S t        dt        |       d       |S # 1 sw Y   xY wc c}w )z1Load all .md files from folder and extract words.u   ⚠️ Folder 'z' not found. Using sample text.z*.mdu   📖 Loading: rzutf-8ignore)encodingerrorsN[a-zA-Z]{2,}r   u   ✅ Extracted z words from markdown files.)r   existsprintsample_wordsglobnameopenreadrefindallextendlowerlen)folder_pathwordsfoldermd_filefcontent	extractedws           r-   load_markdown_filesrH   7   s    E+F==?},KLM~;;v& 5w||n-.'3B 	affhG	 JJ8	3Aaggi345 5zQ~	N3u:,&A
BCL	 	 4s   2D0D
D	c                  N    d} t        j                  d| j                               S )Nz
    intelligence artificial machine learning neural network deep 
    consciousness awareness thinking reasoning planning algorithm 
    r3   )r;   r<   r>   )sample_texts    r-   r6   r6   N   s%    K ::o{'8'8':;;r,   c                    t        t              }| D ]  }||xx   dz  cc<    t        |j                         d       }|d| D cg c]  \  }}|	 }}}t	        |      D ci c]  \  }}||
 }	}}|	j                         D ci c]  \  }}||
 }
}}t        dt        |              |	|
|fS c c}}w c c}}w c c}}w )zBuild word-to-index vocabulary.   c                     | d    S )NrL   r+   )xs    r-   <lambda>z"build_vocabulary.<locals>.<lambda>[   s    adU r,   )keyNu   📚 Vocabulary size: )r   intsorteditems	enumerater5   r?   )rA   max_vocab_sizeword_countswordsorted_wordsrG   cvocab_wordsidxword2idxidx2words              r-   build_vocabularyr^   U   s    c"K DQ +++-?CL!-o~!>?A1?K?+4[+ABic4c	BHB+3>>+;<idCT	<H<	"3{#3"4
56X{** @B<s   B4%B:C c                    t        j                  d||fd      }t        j                  |      }	 t	        j
                  d|      }|j                  d| |      }|d   |d	   z
  }|d
   |d   z
  }||z
  dz  }	||z
  dz  }
|j                  |	|
f| d	|       t        j                  |t        j                        dz  }d|z
  }|S #  	 t	        j
                  d|      }n#  t	        j                         }Y nxY wY xY w)z)Render a word as a grayscale 28x28 image.L   )colorz//usr/share/fonts/truetype/dejavu/DejaVuSans.ttfz	arial.ttf)r   r   )font   r   r   rL   )fillrc   dtypeg     o@      ?)r   newr	   Drawr
   truetypeload_defaulttextbboxtextnparrayfloat32)rW   sizer   imgdrawrc   bbox
text_widthtext_heightrN   y	img_arrays               r-   render_word_imagerz   h   s	   
))C$S
1C>>#D,!!"SU^_ ==D=1Da47"Jq'DG#K	
	q A		!AIIq!fdI.BJJ/%7IiI%,	,%%k9=D	,))+Ds#   C C>
C! C>!C97C>c                   $    e Zd ZdZddZd Zd Zy)SequenceDatasetz:Dataset of word image sequences with next-word prediction.c                    g | _         || _        || _        t        |      | _        t        |j                               }|D cg c]	  }||v s| }}t        |      |kD  r|d | }|| _        t        t        |      |z
        D ]  }||||z    }	|||z      }
g }g }|	D ]F  }	 t        |      }|j                  t        j                  |             |j                  ||          H ||
   }| j                   j                  |||d        t!        dt        | j                          d| d       y c c}w #  t        j                  dt        j                        }Y xY w)N)r   r   rf   )imagesinput_indices
target_idxu   🖼️ Sequence dataset: z samples (seq_len=))samplesr\   r"   r?   r!   setkeysrA   rangerz   ro   zerosrq   appendr(   FloatTensorr5   )selfrA   r\   r"   max_samples	vocab_setrG   valid_wordsiinput_wordstarget_wordr~   indicesrs   r   s                  r-   __init__zSequenceDataset.__init__   ss    x=(	"':Q1	>q:: {k)%l{3K 
 s;''12 	A%aG4K%a'k2K FG  ,?+A.C e//45x{+, "+.JLL !((! #	. 	*3t||+<*==OPWyXYZ[A ;&?((82::>Cs   	D2D2!D77'E c                 ,    t        | j                        S )N)r?   r   )r   s    r-   __len__zSequenceDataset.__len__   s    4<<  r,   c                     | j                   |   }t        j                  |d         j                  d      }t        j                  |d         }t        j                  |d   g      }|||fS )Nr~   rL   r   r   )r   r(   stack	unsqueeze
LongTensor)r   r[   sampler~   r   r   s         r-   __getitem__zSequenceDataset.__getitem__   sh    c"VH-.88;(()@A%%vl';&<=
}j00r,   N)r   iP  )r   r   r   __doc__r   r   r   r+   r,   r-   r|   r|      s    D'\R!1r,   r|   c                     t        |  \  }}}t        j                  |      }t        j                  |      }t        j                  |      j                  d      }|||fS )zPad sequences to same length.)zipr(   r   squeeze)batchimages_listinput_indices_listtarget_idx_listr~   r   r   s          r-   collate_seq_batchr      s\    7:E{4K#_ [[%F KK 23M _-55b9J=*,,r,   c                   *     e Zd ZdZd fd	Zd Z xZS )WordMLPEncoderz'Conv2D + Dense encoder for word images.c                 <   t         |           t        j                  t        j                  dddd      t        j
                         t        j                  d      t        j                  dddd      t        j
                         t        j                  d      t        j                  dddd      t        j
                         t        j                  d      	      | _        t        j                         | _	        t        j                  d	|      | _        y )
NrL       r   )kernel_sizepaddingrd   )r   r   r   i  )superr   nn
SequentialConv2dReLU	MaxPool2dfeaturesFlattenflattenLinear	embedding)r   r   	__class__s     r-   r   zWordMLPEncoder.__init__   s    IIaA6GGILLQ'IIb"!Q7GGILLQ'IIb#1a8GGILLQ'
 zz|;	:r,   c                    |j                  d      |j                  d      }}|j                  ||z  ddd      }| j                  |      }| j                  |      }| j	                  |      }|j                  ||d      }|S )zx: [batch, seq_len, 1, 28, 28]r   rL   r   r   )rr   viewr   r   r   )r   rN   r#   r"   s       r-   forwardzWordMLPEncoder.forward   sz    ffQiG
 FF:'B3 MM!LLONN1 FF:w+r,   )r   )r   r   r   r   r   r   __classcell__r   s   @r-   r   r      s    1;(r,   r   c                   >     e Zd ZdZd fd	ZddZd Zd Zd Z xZ	S )	SequencePredictorz
    Sequence model with ODE-CCT dynamics:
    - LSTM for trajectory encoding
    - Entropy-based confidence for early stopping
    - Periodic pattern detection
    c           	         t         |           || _        || _        || _        t        |      | _        t        j                  ||ddd      | _	        t        j                  t        j                  ||      t        j                         t        j                  d      t        j                  ||            | _        t        j                  ||      | _        | j!                  dt#        j$                  d             d| _        y )	Nrd   Tg?)
input_sizehidden_size
num_layersbatch_firstdropoutg?entropy_historyd   r   )r   r   r   r    r!   r   encoderr   LSTMlstmr   r   r   Dropoutdecoder	predictorregister_bufferr(   r   entropy_idxr   r   r    r!   r   s       r-   r   zSequencePredictor.__init__  s    "$& &i0 GG "
	 }}IIj*-GGIJJsOIIj),	
 9k: 	.C0@Ar,   c                 r   | j                  |      }| j                  |      \  }\  }}|d   }| j                  |      }| j                  |      }	|rct	        j
                  |	d      }
|
t        j                  |
dz         z  j                  d       }|	|j                         j                         fS |	|fS )z1Forward pass with optional confidence estimation.r   rL   dimg:0yE>)r   r   r   r   Fsoftmaxr(   logsummeanitem)r   r~   return_confidence
embeddingslstm_outhiddencelllast_hiddendecodedlogitsprobsentropys               r-   r   zSequencePredictor.forward%  s     \\&)
 $(99Z#8 .64 Rj ,,{+ (IIf!,E		%$, 77<<<CCG7<<>..000wr,   c                 z    t        |      t        j                  k  ry|t        j                   d }t        |      S )z:Detect if model is in a periodic (high-frequency) pattern.FN)r?   cfgr'   all)r   predictionstargetsrecent_corrects       r-   detect_periodicityz$SequencePredictor.detect_periodicity=  s;     {c444$c&<&<%<%=>>""r,   c                 j    | j                   dz  }|| j                  |<   | xj                   dz  c_         y)z,Update entropy history for ODE-CCT tracking.r   rL   N)r   r   )r   entropy_valuer[   s      r-   update_entropyz SequencePredictor.update_entropyF  s3    $$1S!Ar,   c                 \    | j                   dd j                         j                         }|S )z+Get stability measure from entropy history.N2   )r   r   r   )r   recents     r-   get_trajectory_stabilityz*SequencePredictor.get_trajectory_stabilityL  s+    %%cr*//1668r,   r   r     )F)
r   r   r   r   r   r   r   r   r   r   r   s   @r-   r   r      s"     D0#r,   r   c                   2     e Zd ZdZd fd	ZddZd Z xZS )ConditionalCollapsePredictorz
    Enhanced predictor with Conditional Collapse Theory.
    - Adaptively allocates compute based on entropy
    - Uses periodicity detection to skip computation
    c                 b    t         |           t        |||      | _        d| _        i | _        y )Nffffff?)r   r   r   
base_modelconfidence_thresholdphrase_cacher   s       r-   r   z%ConditionalCollapsePredictor.__init__Y  s2    +Iz;O %(! r,   c                 :   |j                  d      }| j                  |      }|r6|| j                  v r(| j                  |   }| j                  |      \  }}||dfS | j                  |      \  }}t	        j
                  |d      }|j                  d      d   j                         j                         }	|	| j                  kD  rR|j                  d      }
t        ||
j                               D ]#  \  }}|| j                  vs|| j                  |<   % ||dfS )a  
        Forward with optional CCT collapse.
        
        If use_collapse=True:
        - Check if sequence matches known periodic pattern
        - If high confidence, skip LSTM (use cached prediction)
        - If low confidence, run full model
        r   TrL   r   F)rr   _get_phrase_keyr   r   r   r   maxr   r   r   argmaxr   tolist)r   r~   use_collapser#   
phrase_keycached_predsr   r   r   
confidencepred_idxrP   preds                r-   r   z$ConditionalCollapsePredictor.forwardd  s    [[^
 ))&1
J$*;*;;,,Z8L"oof5OFG7D(( //&1 		&a(YY1Y%a(--/446
111}}}+H X__->? 2	Td///-1D%%c*2 w%%r,   c                 H    t        |ddddddf   j                               S )z7Extract discrete key from image sequence (for caching).Nr      )tupler   )r   r~   s     r-   r   z,ConditionalCollapsePredictor._get_phrase_key  s'     VAq!RO,33566r,   r   )T)r   r   r   r   r   r   r   r   r   s   @r-   r   r   R  s    	$&L7r,   r   c                    t         j                  j                  | j                         |      }t	        j
                         }t         j                  j                  j                  |dd      }| j                  t        j                         t        |      D ]  }| j                          d}	d}
d}d}d}t        |      D ])  \  }\  }}}|j                  t        j                        }|j                  t        j                        }|j                           | |d      \  }} |||      }|j                          t         j                  j                   j#                  | j                         d	
       |j%                          |	|j'                         z  }	|j)                  d      }|
||k(  j+                         j'                         z  }
||j-                  d      z  }||z  }, d|
z  |z  }|	t/        |      z  }|t/        |      z  }| j1                          d}d}d}t        j2                         5  |D ]  \  }}}|j                  t        j                        }|j                  t        j                        } | |d      \  }} |||      }||j'                         z  }|j)                  d      }|||k(  j+                         j'                         z  }||j-                  d      z  } 	 ddd       d|z  |z  }|t/        |      z  }|j%                  |       t5        d|dz    d| d|dd|dd|dd|dd|d        | S # 1 sw Y   ]xY w)z Train sequence prediction model.)r%   r         ?)patiencefactor        r   Tr   rh   )max_normrL   r   r   NzEpoch /z	 | Loss: .4fz
 | Train: .2fz% | Val Loss:  | Val: z% | Entropy: z.3f)r(   optimAdam
parametersr   CrossEntropyLosslr_schedulerReduceLROnPlateautor   r)   r   trainrT   	zero_gradbackwardutilsclip_grad_norm_stepr   r   r   rr   r?   evalno_gradr5   )modeltrain_loader
val_loaderr$   r%   	optimizer	criterion	schedulerepoch
train_losstrain_correcttrain_totalentropy_sumcollapse_count	batch_idxr~   r   r   r   r   losspreds	train_accavg_lossavg_entropyval_lossval_correct	val_total_val_accavg_val_losss                                  r-   train_modelr8    s=      !1!1!3 ;I##%I((::9qY\:]I	HHSZZv @-
>G>U 	#:I:zYYszz*F#szz2J! $FdCOFG VZ0DMMOHHNN**5+;+;+=*LNN $))+%JMMaM(Eez1668==??M:??1--K7"K-	#0 -'+5	L 11!C$55 	

	]]_ 	05? 
01z3::.']]3::6
!&DA	 4DIIK'!, 388:??AAZ__Q//	
0	0 #i/#j/1|$uQwiq )nJyo >',HWSM B%c*, 	-{@-D L/	0 	0s   >B=MM!	c                 v   | j                          t        j                         }d}d}d}g }g }	g }
g }t        j                         5  |D ]  \  }}}|j                  t        j                        }|j                  t        j                        } | |d      \  }} |||      }||j                         z  }|j                  d      }|||k(  j                         j                         z  }||j                  d      z  }t        j                  |d      }|j                  d      d   }t        |t               r|j#                  |gt%        |      z         n|j#                  |j'                                |	j#                  |j'                                t)        t%        |            D ]:  }||   j                         ||   j                         k(  }|j+                  |       <  	 ddd       d|z  |z  }t-        j.                  |      }t-        j.                  |	      }t-        j0                  |      }t3        d	       t3        d
       t3        d       t3        d|t%        |      z  d       t3        d|dd| d| d       t3        d|d       t3        d|d       t3        d|d       t3        d       dddddddt!        d      g}dgt%        |      dz
  z  }|D ]D  }t)        t%        |      dz
        D ](  }||   |cxk  r||dz      k  sn ||xx   dz  cc<    D F t3        d       g d}t5        ||      D ]>  \  }} d| z  t%        |      z  }!d t7        |!d!z        z  }"t3        d"|d#d$|" d%|!d&d'       @ t        d( |	D              }#t        d) |	D              }$t        d* |	D              }%t3        d+       t3        d,d|#z  t%        |	      z  d&d'       t3        d-d|$z  t%        |	      z  d&d'       t3        d.d|%z  t%        |	      z  d&d'       t3        d/       | j                          d}&t        j                         5  |D ]  \  }}}t)        t9        d0t%        |                  D ]  }||   D 'cg c]"  }'|j;                  |'j                         d1      $ }(}'|j;                  ||   j                         d1      })|||dz    }* | |*d      \  }}+|j                  d      j                         },|j;                  |,d1      }-|,||   j                         k(  rd2nd3}.t3        d4|. d5d%j=                  |(d6d        d7|) d8|- d9	       |&dz  }&|&d:k\  s n |&d:k\  s n ddd       t3        d       |||fS # 1 sw Y   xY wc c}'w # 1 sw Y   +xY w);z"Test and analyze ODE-CCT dynamics.r
  r   Tr  rL   r   Nr   z=
============================================================u%   📊 SEQUENCE PREDICTION TEST RESULTS<============================================================zTest Loss: r  zTest Accuracy: r  z% (r  r   zAverage Entropy: zAverage Confidence: Perplexity: u    
🧠 ODE-CCT DYNAMICS ANALYSIS:r  rh   g      ?g       @g      @g      @infz1
  Entropy Distribution (lower = more confident):)z<0.5z0.5-1.0z1.0-1.5z1.5-2.0z2.0-3.0z3.0-5.0z>5.0u   █rd   z    z>8z:  z.1f%c              3   ,   K   | ]  }|d kD  s	d  yw)?rL   Nr+   .0rY   s     r-   	<genexpr>ztest_model.<locals>.<genexpr>9  s     <!AGA<   
c              3   >   K   | ]  }d |cxk  rdk  sn nd  yw)r  r@  rL   Nr+   rA  s     r-   rC  ztest_model.<locals>.<genexpr>:  s     B3>c>1Bs   
c              3   ,   K   | ]  }|d k  s	d  yw)r  rL   Nr+   rA  s     r-   rC  ztest_model.<locals>.<genexpr>;  s     <181<rD  z
  Confidence Levels:z    High (>0.9): z    Medium (0.5-0.9): z    Low (<0.5): u   
🔍 Sample Predictions:r   ?u   ✅u   ❌z  z Input: u    → | True: 'z' | Pred: ''r   )r  r   r  r(   r  r  r   r)   r   r   r   rr   r   r   r   
isinstancefloatr=   r?   r   r   r   ro   r   expr5   r   rQ   mingetjoin)/r   test_loaderr\   r]   r$  	test_losstest_correct
test_total	entropiesconfidence_scorescorrect_sequencesall_predictionsr~   r   r   r   r   r-  r.  r   r   r   
is_correcttest_accr1  avg_confidence
perplexityentropy_binsentropy_countselabelslabelcountpctbar	high_confmed_conflow_confsample_countr[   r   	true_word	img_batchr5  r  	pred_wordstatuss/                                                  r-   
test_modelrl    s    
JJL##%IILJ I O	 31< 	3-FM:YYszz*F#szz2J#FdCOFGVZ0D$IMMaM(EUj0557<<>>L*//!,,J IIf!,Eq)!,J'5)  'S_!<=  !12$$Z%6%6%89 3u:& 3"1X]]_
10B0B0DD
&&z23/	338 \!J.H'')$KWW./N $J	-	
12	&M	K	#k"2237
89	OHS>\N!J<q
IJ	k#.
/0	  4
56	LC(
)* 

-. sCc3U5\BLSC-12N s<(1,- 	AA!7l1Q3&77q!Q&!	 

>?TFFN3 3uEkC	N*c#'l"U2JbQs3iq123 <0<<IB/BBH</<<H	"$	c)mC0A,BB3Gq
IJ	"3x<4E0F#Fs"K1
MN	S\#.?*@@EQ
GH 

&'	JJLL	 1< 	-FM:3q#j/23 HUVWHXYx||CHHJ<YY$LLA););)=sC	"1QqSM	!)tD	!==Q=/446$LL37	"*jm.@.@.B"B6((388K4D+E*FnU^T__jktjuuvwx!1$  q %	* 
&M[*,,M3 3d Z s8   F)V#3V/'V*
=B:V/8V/V/V'*V//V8c            
         t        d       t        d       t        d       t        t        j                        } t        d       t	        | d      \  }}}t        |      t        _        t        dt        j                   d       t        | |t        j                  d	
      }t        d       t        |      }t        d|z        }t        d|z        }||z
  |z
  }t        j                  j                  j                  ||||g      \  }	}
}t        |	t        j                  dt               }t        |
t        j                  dt               }t        |t        j                  dt               }t        dt        |	       dt        |
       dt        |              t        d       t#        t        j$                  t        j&                  t        j                        }t        |       t        d       t)        |||t        j*                  t        j,                        }t        d       t/        ||||      \  }}}t        j0                  |j3                         ||t        j$                  t        j&                  t        j                  t        j                  ddd       t        d       t        d|dd       t        d |d       y )!Nu:   🧠 Word-MLP Sequence Model (ODE-CCT Language Prediction)r:  u'   
📂 Step 1: Loading markdown files...u$   
📚 Step 2: Building vocabulary...r   )rU   u4   
🖼️ Step 3: Creating sequence dataset (seq_len=z)...i@  )r"   r   u   
📊 Step 4: Splitting data...r   g333333?T)r#   shuffle
collate_fnFz	  Train: r  z	 | Test: u"   
🏗️ Step 5: Building model...)r   r    r!   u(   
🚀 Step 6: Training sequence model...)r$   r%   u   
🧪 Step 7: Testing...)r   r    r!   r"   )model_state_dictr\   r]   configzword_mlp_sequence.ptu+   
💾 Model saved to 'word_mlp_sequence.pt'zFinal Test Accuracy: r  r>  r;  )r5   rH   r   r   r^   r?   r!   r"   r|   rQ   r(   r  datarandom_splitr   r#   r   r   r   r    r8  r$   r%   rl  save
state_dict)rA   r\   r]   rZ   datasettotal
train_sizeval_size	test_sizetrain_datasetval_datasettest_datasetr!  r"  rP  r   rY  r1  r[  s                      r-   mainr~  d  sU   	
FG	&M 

450E 

12&6uT&R#Hh+&CO 
A#++d
STeXs{{PUVG 

,-LES5[!J4%< H
"X-I/4{{/?/?/L/L*h	20,M; mPTarsLKCNNE^opJ\cnne`qrK	Ic-()#k2B1C9SQ]M^L_
`a 

/0-->>OOE
 
%L 

56|Z

svvVE 

%&(25+xQY(Z%Hk: 
JJ!,,...??{{	
	
 
 

89	!(3q
12	LC(
)*r,   __main__)i  )r   r   )r   r   )+r   osr;   randomnumpyro   r(   torch.nnr   torch.nn.functional
functionalr   torch.utils.datar   r   torch.nn.utils.rnnr   r   r   PILr   r	   r
   collectionsr   pathlibr   r   r   rH   r6   r^   rz   r|   r   Moduler   r   r   r8  rl  r~  r   r+   r,   r-   <module>r     s   
 
 	       0 V V + + # J J4 h.<+&>51g 51p-&&RYY &RT		 Tn<7299 <7FKbx-|B+H zF r,   