Update README.md

README.md

@@ -283,7 +283,6 @@ python main_menu.py --task advanced_inference --query "Your complex query here"
 3. **Thought Sequence Generation:** Produces a sequence of interconnected thoughts/actions.
 4. **Final Response Generation:** Synthesizes the best thought path into a coherent response.
 
-
 ---
 
 ### **Mode: With World Model and Tree of Thought**
@@ -291,162 +290,97 @@ python main_menu.py --task advanced_inference --query "Your complex query here"
 #### **Step 1: Input Tokenization and Encoding**
 
 1. **Tokenization:**
-   \[
-   \text{input\_ids} = \text{tokenizer.encode(query, return\_tensors='pt')}
-   \]
-   - **Shape:** \((\text{batch\_size}=1, \text{seq\_len})\)
+   - The input query is converted into token IDs using the tokenizer. This numerical representation is essential for processing by the Transformer model.
+   - **Shape:** The resulting tensor has a shape corresponding to the batch size and the sequence length of the input.
 
 2. **Encoding via Transformer:**
-   \[
-   \text{transformer\_output} = \text{model\_transformer}(input\_ids, input\_ids)
-   \]
-   - **Shape:** \((\text{batch\_size}=1, \text{seq\_len}, \text{d\_model})\)
+   - The tokenized input is passed through the Transformer model to generate contextual embeddings. These embeddings capture the semantic information of the input.
+   - **Shape:** The output tensor includes the batch size, sequence length, and the dimensionality of the model.
 
 3. **State Representation:**
-   \[
-   \text{initial\_representation} = \text{RepresentationNetwork}(\text{transformer\_output})[:, -1, :].unsqueeze(1)
-   \]
-   - **Shape:** \((\text{batch\_size}=1, 1, \text{state\_dim})\)
+   - The `RepresentationNetwork` processes the transformer's output to create a condensed state representation. This state serves as the foundation for further reasoning steps.
+   - **Shape:** The state representation tensor includes the batch size, a single sequence length (typically one), and the state dimensionality.
 
 4. **State Initialization:**
-   \[
-   \text{initial\_state} = \text{State}(\text{representation}=\text{initial\_representation}, \text{dynamics\_network}=dynamics\_network, \text{action\_encoder}=action\_encoder, \text{thought\_node}=root\_thought\_node)
-   \]
+   - A `State` object is created, encapsulating the initial representation, dynamics network, action encoder, and the root node of the Tree of Thought. This object maintains the current context and facilitates state transitions as actions are applied.
 
 #### **Step 2: MCTS Initialization and Root Node Evaluation**
 
 1. **MCTS Instance Creation:**
-   \[
-   \text{mcts} = \text{MCTS}(\text{prediction\_network}, \text{dynamics\_network}, \text{action\_encoder}, \text{num\_iterations}=mcts\_iterations, \text{exploration\_constant}=exploration\_constant)
-   \]
+   - An instance of the `MCTS` class is initialized with the necessary networks and parameters. This instance will manage the search process through the Tree of Thought.
+   - Key parameters include the prediction network, dynamics network, action encoder, number of iterations, and the exploration constant.
 
 2. **Root Node Creation:**
-   \[
-   \text{root\_node} = \text{MCTSNode}(\text{state}=\text{initial\_state}, \text{thought\_node}=root\_thought\_node)
-   \]
+   - A `MCTSNode` representing the root of the search tree is created. This node is associated with the initial state and the root thought node from the Tree of Thought.
 
 3. **Root Node Evaluation:**
-   \[
-   \text{value\_estimate} = \text{mcts.evaluate}(\text{root\_node})
-   \]
+   - The root node is evaluated using the MCTS's evaluation function, which assesses the potential value of the current state based on the prediction network's output.
 
 4. **Backpropagation:**
-   \[
-   \text{mcts.backpropagate}(\text{root\_node}, \text{value\_estimate})
-   \]
+   - The evaluation result (value estimate) is backpropagated through the tree, updating the visit counts and value sums of the nodes. This process informs future selections by providing aggregated value information.
 
 #### **Step 3: MCTS Iterations with Beam Search**
 
 1. **Beam Initialization:**
-   \[
-   \text{beam} = \left\{ \left( \text{root\_node}, \text{score} = 0, \text{cum\_entropy} = 0, \text{cum\_variance} = 0, \text{action\_sequence} = [] \right) \right\}
-   \]
+   - The search beam is initialized with the root node. Each beam element includes the current node, a cumulative score, cumulative entropy, cumulative variance, and an empty action sequence.
 
 2. **Iterative Expansion:**
-   - **For each iteration up to \(\text{num\_iterations}\):**
+   - **For each iteration up to the specified number of MCTS iterations:**
      
      - **Candidate Collection:**
        - **For each node in the current beam:**
          
          - **Leaf Evaluation:**
-           - **If** \(\text{node.is\_leaf()} = \text{True}\):
-             \[
-             \text{value\_estimate} = \text{mcts.evaluate}(\text{node})
-             \]
-             \[
-             \text{mcts.backpropagate}(\text{node}, \text{value\_estimate})
-             \]
+           - If the node is a leaf (has no children), it is evaluated to estimate its value. The evaluation result is then backpropagated to update the node's statistics.
          
          - **Child Selection:**
-           - **If** \(\text{node.children}\) **is not empty:**
-             - **Calculate total visits:**
-               \[
-               \text{total\_visits} = \sum_{\text{child} \in \text{node.children}} \text{child.visit\_count}
-               \]
-             
-             - **Select top \(\text{beam\_size}\) actions based on UCB scores:**
-               \[
-               \text{sorted\_children} = \text{sorted}(\text{node.children.items()}, \text{key}=\lambda \text{item}: \text{item}[1].ucb\_score(\text{total\_visits}, \text{exploration\_constant}), \text{reverse}=True)[:\text{beam\_size}]
-               \]
+           - If the node has children, the total number of visits to all its children is calculated. The children are then sorted based on their Upper Confidence Bound (UCB) scores, which consider exploration and exploitation factors. The top actions (up to the beam size) are selected for expansion.
+     
+     - **Action Sequence Prediction:**
+       - **For each selected action:**
          
-         - **Action Sequence Prediction:**
-           - **For each selected action:**
+         - **Initialization:**
+           - Set the current node to the selected child node.
+           - Initialize the current action sequence with the selected action.
+           - Initialize the current score, cumulative entropy, and cumulative variance to zero.
+         
+         - **Multi-Token Prediction:**
+           - **For each step in the number of tokens to predict:**
+             
+             - **Leaf Evaluation:**
+               - If the current node is a leaf, evaluate it to obtain a value estimate and backpropagate the result.
+             
+             - **Child Check:**
+               - If the current node has no children, exit the multi-token prediction loop for this sequence.
              
-             - **Initialize:**
-               - \(\text{current\_node} = \text{selected\_node}\)
-               - \(\text{current\_sequence} = [\text{selected\_action}]\)
-               - \(\text{current\_score} = 0\)
-               - \(\text{current\_entropy} = 0\)
-               - \(\text{current\_variance} = 0\)
+             - **Action Selection:**
+               - Calculate the total number of visits to all children of the current node.
+               - Select the action with the highest UCB score from the children.
              
-             - **Multi-Token Prediction:**
-               - **For each step in \(\text{n\_tokens\_predict}\):**
-                 
-                 - **If** \(\text{current\_node.is\_leaf()} = \text{True}\):
-                   \[
-                   \text{value\_estimate} = \text{mcts.evaluate}(\text{current\_node})
-                   \]
-                   \[
-                   \text{mcts.backpropagate}(\text{current\_node}, \text{value\_estimate})
-                   \]
-                 
-                 - **If** \(\text{current\_node.children}\) **is empty**, **break**.
-                 
-                 - **Calculate total visits for the new node:**
-                   \[
-                   \text{total\_visits} = \sum_{\text{child} \in \text{current\_node.children}} \text{child.visit\_count}
-                   \]
-                 
-                 - **Select the action with the highest UCB score:**
-                   \[
-                   (\text{next\_action}, \text{next\_node}) = \text{max}(\text{current\_node.children.items()}, \text{key}=\lambda \text{item}: \text{item}[1].ucb\_score(\text{total\_visits}, \text{exploration\_constant}))
-                   \]
-                 
-                 - **Score Update:**
-                   \[
-                   \text{current\_score} += 
-                   \begin{cases} 
-                   \frac{\text{next\_node.value\_sum}}{\text{next\_node.visit\_count}} & \text{if } \text{next\_node.visit\_count} > 0 \\
-                   0 & \text{otherwise}
-                   \end{cases}
-                   \]
-                 
-                 - **Entropy and Variance Update:**
-                   \[
-                   \text{current\_entropy} += \text{next\_node.entropy}
-                   \]
-                   \[
-                   \text{current\_variance} += \text{next\_node.variance}
-                   \]
-                 
-                 - **Append next action to the sequence:**
-                   \[
-                   \text{current\_sequence}.append(\text{next\_action})
-                   \]
-                 
-                 - **Update current node:**
-                   \[
-                   \text{current\_node} = \text{next\_node}
-                   \]
+             - **Score Update:**
+               - If the selected child node has been visited before, increment the current score by the average value estimate of that node.
+               - If the child node has not been visited, the score remains unchanged or is updated with a default value.
              
-             - **Candidate Aggregation:**
-               \[
-               \text{all\_candidates.append}((\text{current\_node}, \text{current\_score}, \text{current\_entropy}, \text{current\_variance}, \text{current\_sequence}))
-               \]
+             - **Entropy and Variance Update:**
+               - Accumulate the entropy and variance metrics from the selected node to guide the search towards more confident and diverse actions.
+             
+             - **Sequence Extension:**
+               - Append the selected action to the current action sequence.
+               - Update the current node to the selected child node.
+         
+         - **Candidate Aggregation:**
+           - Add the new candidate sequence, along with its updated score, entropy, variance, and action sequence, to the list of all candidates for this iteration.
      
      - **Beam Pruning:**
-       \[
-       \text{beam} = \text{sorted}(\text{all\_candidates}, \text{key}=\lambda x: x[1] - 0.1 \times x[2] + 0.05 \times x[3], \text{reverse}=True)[:\text{beam\_size}]
-       \]
+       - After collecting all candidates from the current beam, sort them based on a scoring function that balances the cumulative score, entropy, and variance.
+       - Retain only the top sequences up to the specified beam size to form the new beam for the next iteration.
 
 3. **Termination:**
-   - **Stop early** if no candidates remain or all beams have reached terminal nodes.
-
+   - The iterative expansion process continues until the specified number of MCTS iterations is reached or there are no more candidates to explore.
+   
 4. **Result Extraction:**
-   \[
-   \text{best\_sequence} = \text{beam}[0][4]
-   \]
-   - **Return:** \(\text{best\_sequence}\) as the generated sequence of actions (thoughts).
+   - After completing the iterations, select the best action sequence from the final beam based on the accumulated scores and metrics.
+   - Return this sequence as the generated series of actions (thoughts) in response to the input query.
 
 ---
 
@@ -455,8 +389,8 @@ python main_menu.py --task advanced_inference --query "Your complex query here"
 #### **1. Step 1: Input Tokenization and Encoding**
 
 - **Tokenization:** The input query is converted into token IDs using the tokenizer. This numerical representation is essential for processing by the Transformer model.
-  
-- **Encoding via Transformer:** The tokenized input is passed through the Transformer model to generate contextual embeddings (`transformer_output`). These embeddings capture the semantic information of the input.
+
+- **Encoding via Transformer:** The tokenized input is passed through the Transformer model to generate contextual embeddings. These embeddings capture the semantic information of the input.
 
 - **State Representation:** The `RepresentationNetwork` processes the transformer's output to create a condensed state representation. This state serves as the foundation for further reasoning steps.
 
@@ -468,7 +402,9 @@ python main_menu.py --task advanced_inference --query "Your complex query here"
 
 - **Root Node Creation:** A `MCTSNode` representing the root of the search tree is created, associated with the initial state and the root thought node.
 
-- **Root Node Evaluation:** The root node is evaluated using the MCTS's evaluation function, which assesses the potential value of the current state. The result is then backpropagated to update the node's statistics.
+- **Root Node Evaluation:** The root node is evaluated using the MCTS's evaluation function, which assesses the potential value of the current state based on the prediction network's output.
+
+- **Backpropagation:** The evaluation result (value estimate) is backpropagated through the tree, updating the visit counts and value sums of the nodes. This process informs future selections by providing aggregated value information.
 
 #### **3. Step 3: MCTS Iterations with Beam Search**
 
@@ -492,7 +428,7 @@ python main_menu.py --task advanced_inference --query "Your complex query here"
   
   - **Candidate Aggregation:** All potential candidates resulting from the action predictions are collected for the current iteration.
 
-- **Beam Pruning:** After all candidates are collected, the beam is pruned to retain only the top sequences based on a scoring function that balances score, entropy, and variance. This ensures that only the most promising action sequences are retained for further exploration.
+- **Beam Pruning:** After collecting all candidates, the beam is pruned to retain only the top sequences based on a scoring function that balances score, entropy, and variance. This ensures that only the most promising action sequences are retained for further exploration.
 
 - **Termination:** The iterative process continues until the specified number of MCTS iterations is reached or all beams have been exhausted.
 
@@ -571,24 +507,6 @@ graph TD
     R --> Q
 ```
 
----
-
-### **Summary**
-
-Your advanced inference mechanism integrates several sophisticated components to enable strategic and coherent response generation:
-
-1. **World Model:** Provides a structured understanding of the current state and predicts future states based on actions.
-2. **Tree of Thought (ToT):** Offers a hierarchical framework for exploring diverse reasoning paths.
-3. **Monte Carlo Tree Search (MCTS):** Efficiently navigates the ToT by balancing exploration and exploitation using UCB scores influenced by entropy and variance.
-4. **Beam Search with Multi-Token Prediction:** Enhances generation efficiency and coherence by predicting multiple tokens simultaneously and maintaining multiple candidate sequences.
-5. **Entropy and Variance:** Quantify uncertainty and diversity, guiding the search process to generate balanced and robust responses.
-
-This comprehensive system enables the model to handle complex queries by simulating strategic reasoning, exploring various thought pathways, and generating informed and coherent responses.
-
----
-
-Feel free to integrate this Markdown into your HuggingFace documentation or repository. If you need further customization or additional sections, let me know!
-
 ## General Arguments
 
 | Argument           | Required | Description                                                                                      | Default             |