Upload 2 files

chain_of_thought_gui.py

@@ -5,14 +5,15 @@ NeuroReasoner Chain-of-Thought GUI (Dark Theme Enhanced)
 A premium Streamlit app for step-by-step reasoning
 across any Hugging Face model (causal or seq2seq).
 Featuring a dark theme, model-type detection, self-consistency
-sampling, and robust handling.
+sampling & voting, robust handling, and GPU telemetry.
 """
 import os
 import time
+import re # Needed for answer normalization
 import streamlit as st
 import torch
 import pynvml # For GPU telemetry
-import numpy as np
+import numpy as np # Imported, but currently unused in core logic
 from transformers import (
     AutoConfig,
     AutoTokenizer,
@@ -23,58 +24,66 @@ from transformers import (
 )
 from collections import Counter # For self-consistency voting
 import gc # Import garbage collector
-
-# Assuming chain_of_thought_wrapper.py is in the same directory
-# and is designed to work with standard Hugging Face models and GenerationConfig.
-# Make sure the wrapper correctly handles num_return_sequences for CoT and SC,
-# and returns the expected dictionary structure:
-# {'full_texts': [...], 'reasoning_steps': [...], 'final_answers': [...], 'consensus_answer': '...'}
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+# --- Import the Enhanced ChainOfThoughtWrapper ---
+# Assuming chain_of_thought_wrapper.py is in the same directory.
+# The wrapper is expected to:
+# 1. Accept model, tokenizer, GenerationConfig (base), device, etc. in __init__.
+# 2. Have a .generate() method that takes input_text (str), GenerationConfig (overrides),
+#    and crucially, num_return_sequences (int) to generate multiple chains efficiently.
+# 3. Return a dictionary with keys 'full_texts', 'reasoning_steps', 'final_answers' (lists).
 try:
     from chain_of_thought_wrapper import ChainOfThoughtWrapper
 except ImportError:
-    st.error("Error: chain_of_thought_wrapper.py not found. Please ensure it's in the same directory.")
-    st.stop()
+    st.error("Error: chain_of_thought_wrapper.py not found. Please ensure the enhanced wrapper script is in the same directory.")
+    st.stop() # Halt execution if the wrapper is not found
 
+# --- Logging Setup for GUI ---
+# Use Streamlit's built-in logging or configure a separate logger
+# For this example, we'll keep it simple and rely mostly on st.status and st.exception
+# if needed, a more detailed logger could be configured here.
 
 # --- Page Configuration ---
 st.set_page_config(
     page_title="🧠 NeuroReasoner CoT GUI",
     page_icon="🧠",
-    layout="wide",
-    initial_sidebar_state="expanded",
+    layout="wide", # Use wide layout
+    initial_sidebar_state="expanded", # Sidebar open by default
     menu_items={
-        'Get Help': 'https://github.com/your_repo_link_here', # Replace or remove
-        'Report a bug': "https://github.com/your_repo_link_here/issues", # Replace or remove
+        'Get Help': 'https://github.com/ayjays132/NeuroReasoner', # Example repo link
+        'Report a bug': "https://github.com/ayjays132/NeuroReasoner/issues", # Example repo issues link
         'About': """
         **NeuroReasoner Chain-of-Thought GUI**
-        An open-source interface powered by Hugging Face models and the NeuroReasoner wrapper.
+        An open-source interface powered by Hugging Face models and the enhanced NeuroReasoner wrapper.
         Explore step-by-step reasoning with various language models.
+        \n\n**Features:** Dark Theme, GPU Telemetry, Model Caching, Self-Consistency Voting,
+        Robust Generation Parameters, Support for Causal and Seq2Seq models.
         """
     }
 )
 
 # --- Dark Theme CSS ---
+# Comprehensive CSS for a professional dark theme inspired by VS Code.
 st.markdown("""
 <style>
     /* Overall Page Background & Text (Dark Theme) */
-    body {
+    /* Target the main container and the root app div */
+    .stApp {
         background-color: #1E1E1E; /* Dark grey background */
         color: #D4D4D4; /* Light grey text */
         font-family: 'Segoe UI', Roboto, Arial, sans-serif;
     }
-    .stApp {
-        background-color: #1E1E1E;
-        color: #D4D4D4;
-    }
 
     /* Sidebar Styling */
     .stSidebar {
         background-color: #2D2D2D; /* Slightly lighter dark grey for sidebar */
         padding: 2rem 1rem;
         border-right: 1px solid #3E3E3E; /* Subtle border */
+        color: #D4D4D4; /* Ensure text in sidebar is light */
     }
     .stSidebar h1, .stSidebar h2, .stSidebar h3 {
-         color: #569CD6; /* Visual Studio Code blue for sidebar headers */
+          color: #569CD6 !important; /* Visual Studio Code blue for sidebar headers */
     }
     .stSidebar label {
         color: #D4D4D4 !important; /* Ensure sidebar labels are visible */
@@ -82,9 +91,7 @@ st.markdown("""
 
 
     /* Main Content Area */
-    .stContainer {
-        padding: 2rem;
-    }
+    /* No specific background needed here, .stApp covers it */
 
     /* Titles and Headers */
     h1, h2, h3, h4, h5, h6 {
@@ -107,6 +114,7 @@ st.markdown("""
         font-weight: bold;
         transition: background-color 0.2s ease, transform 0.1s ease;
         box-shadow: 2px 2px 5px rgba(0, 0, 0, 0.3);
+        margin-top: 1.65rem; /* Add top margin to align with text area */
     }
     .stButton>button:hover {
         background-color: #27633A; /* Lighter green on hover */
@@ -120,7 +128,9 @@ st.markdown("""
 
 
     /* Text areas and inputs */
-    .stTextArea textarea, .stTextInput input {
+    /* Target specific classes used by Streamlit for input/text areas */
+    div[data-baseweb="textarea"] textarea,
+    div[data-baseweb="input"] input {
         border: 1px solid #3E3E3E; /* Dark border */
         border-radius: 0.4rem;
         padding: 0.75rem;
@@ -129,26 +139,28 @@ st.markdown("""
         color: #D4D4D4; /* Light text */
         box-shadow: inset 1px 1px 3px rgba(0, 0, 0, 0.2);
     }
-     .stTextArea label, .stTextInput label {
+     div[data-baseweb="textarea"] label,
+     div[data-baseweb="input"] label,
+     .stSlider label, .stSelectbox label, .stCheckbox label {
         font-weight: bold;
         color: #9CDCFE !important; /* Light blue labels */
         margin-bottom: 0.5rem;
         display: block;
     }
-    /* Streamlit status box styling */
-    .st-emotion-cache-vj1l9j { /* Target the status box content div */
+     /* Streamlit status box styling - Target the main container and its contents */
+    .st-emotion-cache-vj1l9j { /* This class might change with Streamlit versions */
         background-color: #2D2D2D; /* Match sidebar background */
         border: 1px solid #3E3E3E;
         border-radius: 0.5rem;
         padding: 1rem;
         margin-bottom: 1rem;
     }
-     .st-emotion-cache-vj1l9j .stMarkdown p { /* Style text inside status */
-        color: #D4D4D4 !important;
-     }
-     /* Status box icons/text (might need to target specific internal classes) */
-     .st-emotion-cache-vj1l9j .stAlert {
-        background-color: transparent !important; /* Don't want alert backgrounds inside status */
+     .st-emotion-cache-vj1l9j .stMarkdown p,
+     .st-emotion-cache-vj1l9j .stAlert { /* Style text and alerts inside status */
+         color: #D4D4D4 !important;
+         background-color: transparent !important; /* Don't want alert backgrounds inside status */
+         border: none !important; /* No borders for alerts inside status */
+         padding: 0.5rem 0 !important; /* Adjust padding */
      }
 
 
@@ -159,11 +171,12 @@ st.markdown("""
         padding: 1rem;
         font-size: 1rem;
         border-left: 5px solid transparent; /* Base style */
+        color: #D4D4D4; /* Default text color for alerts */
     }
-    .stAlert.stAlert-info { border-left-color: #569CD6; background-color: #2A3E52; color: #9CDCFE; } /* Dark blue info */
-    .stAlert.stAlert-success { border-left-color: #4EC9B0; background-color: #28403A; color: #7AC7A3; } /* Dark teal success */
-    .stAlert.stAlert-warning { border-left-color: #DCDCAA; background-color: #454032; color: #FFDAA6; } /* Dark yellow warning */
-    .stAlert.stAlert-error { border-left-color: #F44747; background-color: #4A3030; color: #F48787; } /* Dark red error */
+    .stAlert.stAlert-info { border-left-color: #569CD6; background-color: #2A3E52; } /* Dark blue info */
+    .stAlert.stAlert-success { border-left-color: #4EC9B0; background-color: #28403A; } /* Dark teal success */
+    .stAlert.stAlert-warning { border-left-color: #DCDCAA; background-color: #454032; } /* Dark yellow warning */
+    .stAlert.stAlert-error { border-left-color: #F44747; background-color: #4A3030; } /* Dark red error */
 
 
     /* Expander styling */
@@ -192,6 +205,10 @@ st.markdown("""
         margin-top: 0;
         color: #D4D4D4;
     }
+    .streamlit-expanderContent .stMarkdown p {
+         color: #D4D4D4 !important; /* Ensure text inside expanders is light */
+    }
+
 
     /* Labels for the output text areas */
     .output-label {
@@ -204,7 +221,8 @@ st.markdown("""
     }
 
     /* Custom class for output text areas to differentiate from input */
-    .output-text-area textarea {
+    /* Need to target the specific Streamlit internal class for the text area */
+    .output-text-area div[data-baseweb="textarea"] textarea {
         background-color: #1E1E1E; /* Even darker background for outputs */
         border: 1px solid #3E3E3E;
         border-radius: 0.4rem;
@@ -238,11 +256,13 @@ st.markdown("""
         font-weight: bold;
     }
     .consensus-answer strong {
-        color: #4EC9B0; /* Teal for "Consensus Answer" label */
-    }
-    .consensus-answer div {
-        color: #D4D4D4; /* Ensure the answer text is light */
+        color: #4EC9B0 !important; /* Teal for "Consensus Answer" label */
     }
+     .consensus-answer p {
+         color: #D4D4D4 !important; /* Ensure the answer text is light */
+         margin: 0 !important; /* Remove default paragraph margins */
+         padding: 0 !important; /* Remove default paragraph padding */
+     }
 
 
 </style>
@@ -250,35 +270,42 @@ st.markdown("""
 
 
 # --- GPU Telemetry Setup ---
+# Initialize NVML for GPU monitoring if available
 try:
     pynvml.nvmlInit()
     GPU_AVAILABLE = True
+    # Get the number of devices to pick the first one (index 0)
+    GPU_COUNT = pynvml.nvmlDeviceGetCount()
+    if GPU_COUNT == 0:
+        GPU_AVAILABLE = False
+        st.warning("NVML initialized but no NVIDIA GPUs found.")
 except Exception:
     GPU_AVAILABLE = False
+    # st.info("NVIDIA Management Library (pynvml) not found or failed to initialize. GPU telemetry disabled.")
 
 # Use st.empty to hold the telemetry status text, defined *outside* cached functions
 telemetry_placeholder = st.empty()
 
 def update_telemetry():
     """Updates the telemetry display in the dedicated placeholder."""
-    telemetry_text = "[Checking System Status...]"
+    telemetry_text = "📊 System Status: [Initializing...]"
     if not GPU_AVAILABLE or not torch.cuda.is_available():
         telemetry_text = "📊 System Status: [No GPU Available]"
     else:
         try:
-            h = pynvml.nvmlDeviceGetHandleByIndex(0)
-            u = pynvml.nvmlDeviceGetUtilizationRates(h)
-            m = pynvml.nvmlDeviceGetMemoryInfo(h)
-            mem_used_mb = m.used // 1024**2
-            mem_total_mb = m.total // 1024**2
-            telemetry_text = f"📊 System Status: GPU {u.gpu}% | Mem {mem_used_mb}/{mem_total_mb} MB"
+            handle = pynvml.nvmlDeviceGetHandleByIndex(0) # Use the first GPU
+            utilization = pynvml.nvmlDeviceGetUtilizationRates(handle)
+            memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
+            mem_used_mb = memory.used // 1024**2
+            mem_total_mb = memory.total // 1024**2
+            telemetry_text = f"📊 System Status: GPU 0: {utilization.gpu}% | Mem {mem_used_mb}/{mem_total_mb} MB"
         except Exception:
+             # If NVML fails after initialization, report error
              telemetry_text = "📊 System Status: [Telemetry Error]"
 
-    # Use markdown with a custom class for styling
+    # Use markdown with a custom class for styling the container
     telemetry_placeholder.markdown(f'<div class="telemetry-box">{telemetry_text}</div>', unsafe_allow_html=True)
 
-
 # Initial telemetry update when the script starts
 update_telemetry()
 
@@ -286,59 +313,83 @@ update_telemetry()
 # --- Caching Model Loading (Core Logic Only) ---
 # Use st.cache_resource for heavy objects like models and tokenizers.
 # This function MUST NOT call Streamlit elements that affect the layout
-# or state outside of its own scope.
-@st.cache_resource(show_spinner=False) # Spinner handled manually
+# or state outside of its own scope (except for the final return value).
+@st.cache_resource(show_spinner=False) # Spinner handled manually in safe_load_model_with_status
 def _load_model_and_tokenizer_cached(model_name: str, device: str, forced_model_type: str = None):
     """
-    Loads the model and tokenizer. This function is cached and should
-    contain minimal Streamlit calls to avoid caching issues.
+    Loads the model and tokenizer. This function is cached by Streamlit.
+    It should perform resource-intensive loading only.
     """
+    # Use low_cpu_mem_usage=True to reduce RAM usage during loading
     config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, low_cpu_mem_usage=True)
     is_encoder_decoder = getattr(config, "is_encoder_decoder", False)
     detected_type = "Seq2Seq" if is_encoder_decoder else "Causal"
 
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-    # Ensure padding token is set for generation robustness
+
+    # Ensure padding token is set for generation robustness, especially for batching (num_return_sequences)
+    # This mirrors the logic in the wrapper's __init__ but is good to do here too
+    # before the model is potentially loaded with a different vocab size.
     if tokenizer.pad_token is None:
         if tokenizer.eos_token is not None:
              tokenizer.pad_token = tokenizer.eos_token
+             # Set the ID explicitly as well
+             tokenizer.pad_token_id = tokenizer.eos_token_id
+             # logger.warning(f"Tokenizer pad_token is None, using eos_token '{tokenizer.eos_token}' as pad_token.")
         else:
-             # Fallback - adding tokens might require resizing model embeddings
-             # which is complex and model-dependent. This is a basic attempt.
-             tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-             tokenizer.pad_token = '[PAD]' # Set the attribute
-             # Attempt to get the new pad token ID - may not work for all tokenizers
-             try:
-                 tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids('[PAD]')
-             except Exception:
-                 tokenizer.pad_token_id = None # Indicate failure to get ID
+            # Fallback: Add a new pad token if neither eos nor pad exists.
+            # The wrapper's init will handle resizing embeddings if possible.
+            # logger.warning("Tokenizer has no pad_token and no eos_token. Adding a new [PAD] token.")
+            tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+            # Need to get the ID for the new token
+            tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids('[PAD]')
+            # logger.info(f"Added new [PAD] token with ID {tokenizer.pad_token_id}.")
+
 
     # Determine the model class based on detection or forced selection
     actual_model_type = forced_model_type if forced_model_type != "Auto" else detected_type
+    model = None # Initialize model to None
 
-    if actual_model_type == "Seq2Seq":
-         model = AutoModelForSeq2SeqLM.from_pretrained(model_name, config=config, trust_remote_code=True)
-    elif actual_model_type == "Causal":
-         model = AutoModelForCausalLM.from_pretrained(model_name, config=config, trust_remote_code=True)
-    else:
-         raise ValueError(f"Unsupported model type selected: {actual_model_type}. Please select 'Auto', 'Causal', or 'Seq2Seq'.")
+    try:
+        if actual_model_type == "Seq2Seq":
+             model = AutoModelForSeq2SeqLM.from_pretrained(model_name, config=config, trust_remote_code=True)
+        elif actual_model_type == "Causal":
+             model = AutoModelForCausalLM.from_pretrained(model_name, config=config, trust_remote_code=True)
+        else:
+             raise ValueError(f"Unsupported model type selected: {actual_model_type}. Please select 'Auto', 'Causal', or 'Seq2Seq'.")
+
+        # Move model to device and set to eval mode
+        model.to(device)
+        model.eval() # Crucial for consistent inference behavior and disabling dropout etc.
 
-    model.to(device)
-    model.eval() # Crucial for consistent inference behavior and disabling dropout etc.
+        # Ensure return_dict_in_generate is True for structured outputs (needed by wrapper for scores/sequences)
+        if not getattr(model.config, 'return_dict_in_generate', False):
+             model.config.return_dict_in_generate = True
+        # Request output_scores by default for potential future CISC use in the GUI voter
+        if not getattr(model.config, 'output_scores', False):
+             model.config.output_scores = True
 
-    # Ensure return_dict_in_generate is True for structured outputs
-    if not getattr(model.config, 'return_dict_in_generate', False):
-         model.config.return_dict_in_generate = True
+        # The wrapper's __init__ will perform its own pad token handling and embedding resizing check.
+        # We ensure the tokenizer passed to it has a pad_token_id.
+
+    except Exception as e:
+        # Clean up resources if model loading failed
+        if model is not None: del model
+        if tokenizer is not None: del tokenizer
+        if torch.cuda.is_available(): torch.cuda.empty_cache()
+        gc.collect()
+        raise e # Re-raise the exception for the caller to handle status updates
 
     return model, tokenizer, actual_model_type
 
 # --- Wrapper function to handle status reporting for cached loading ---
 def safe_load_model_with_status(model_name: str, device: str, forced_model_type: str = None):
     """
-    Calls the cached loading function and handles Streamlit status updates.
+    Calls the cached loading function (_load_model_and_tokenizer_cached)
+    and handles Streamlit status updates and error reporting.
     """
+    # Use st.status here, defined outside the cached function, for live updates
     status_text = f"🌐 Loading model '{model_name}' on device '{device}'..."
-    # Use st.status here, defined outside the cached function
     with st.status(status_text, expanded=True) as status_box:
         status_box.write("Checking system status...")
         update_telemetry() # Update the separate telemetry box
@@ -352,9 +403,9 @@ def safe_load_model_with_status(model_name: str, device: str, forced_model_type:
                 forced_model_type=forced_model_type
             )
 
-            # Report padding token status if available
+            # Report padding token status after loading
             if tokenizer and tokenizer.pad_token_id is None:
-                 status_box.warning(f"Tokenizer has no pad_token_id. Generation might fail for models requiring padding (e.g., batching).")
+                 status_box.warning(f"Tokenizer has no pad_token_id. Batch generation (Self-Consistency) might be unstable.")
             elif tokenizer:
                  status_box.write(f"Tokenizer pad_token_id set to {tokenizer.pad_token_id}.")
 
@@ -366,21 +417,81 @@ def safe_load_model_with_status(model_name: str, device: str, forced_model_type:
         except Exception as e:
             status_box.error(f"❌ Model loading failed.")
             update_telemetry() # Final telemetry update after error
-            st.exception(e) # Display the full exception traceback
-            # Clean up resources in case of failure before returning None
-            # These are manual attempts; cache handles cleanup on its own state changes
-            # but explicit cleanup is good practice on error paths.
-            try:
-                if 'model' in locals() and model is not None: del model
-            except NameError: pass
-            try:
-                if 'tokenizer' in locals() and tokenizer is not None: del tokenizer
-            except NameError: pass
-            if torch.cuda.is_available(): torch.cuda.empty_cache()
-            gc.collect()
+            st.exception(e) # Display the full exception traceback within the status box
+            # No need for manual cleanup here, as the exception in the cached function
+            # should have triggered cleanup within that function, and Streamlit's
+            # cache resource management handles state on failure.
             return None, None, None # Return None on failure
 
 
+# --- Self-Consistency Voting Logic ---
+def normalize_answer(answer: str) -> str:
+    """
+    Normalizes a string answer for robust comparison during voting.
+    - Converts to lowercase.
+    - Strips leading/trailing whitespace.
+    - Removes common punctuation.
+    - Can be extended with more sophisticated normalization (e.g., number words to digits).
+    """
+    if not isinstance(answer, str):
+        return "" # Handle non-string inputs
+
+    # Simple normalization: lowercase, strip whitespace, remove common punctuation
+    normalized = answer.lower().strip()
+    # Remove common trailing characters like periods, commas, etc.
+    normalized = re.sub(r'[.,!?;:]+$', '', normalized).strip()
+    # Remove common leading "Answer: " or similar preambles (case-insensitive)
+    normalized = re.sub(r'^\s*(?:the answer is|result|output)\s*[:\-]?\s*', '', normalized, flags=re.IGNORECASE).strip()
+    # Add more normalization rules if needed (e.g., handling "forty two" vs "42")
+
+    return normalized
+
+def perform_self_consistency_voting(final_answers: List[str]) -> Tuple[Optional[str], Dict[str, int]]:
+    """
+    Performs simple majority voting on a list of final answers.
+    Filters out empty answers and normalizes them before voting.
+
+    Args:
+        final_answers (List[str]): A list of raw final answer strings from the wrapper.
+
+    Returns:
+        Tuple[Optional[str], Dict[str, int]]: A tuple containing:
+            - The winning (most common) normalized answer, or None if no valid answers.
+            - A dictionary mapping normalized answers to their vote counts.
+    """
+    if not final_answers:
+        return None, {}
+
+    # 1. Filter out empty or non-string answers
+    valid_answers = [ans for ans in final_answers if isinstance(ans, str) and ans.strip()]
+
+    if not valid_answers:
+        return None, {}
+
+    # 2. Normalize answers
+    normalized_answers = [normalize_answer(ans) for ans in valid_answers]
+    # Filter out answers that became empty after normalization
+    normalized_answers = [ans for ans in normalized_answers if ans.strip()]
+
+    if not normalized_answers:
+         return None, {}
+
+
+    # 3. Perform majority voting
+    answer_counts = Counter(normalized_answers)
+
+    # 4. Determine the consensus answer
+    # most_common(1) returns a list like [('answer', count)]
+    most_common_item = answer_counts.most_common(1)
+
+    if most_common_item:
+        consensus_answer = most_common_item[0][0]
+        return consensus_answer, dict(answer_counts)
+    else:
+        # This case should ideally not happen if normalized_answers is not empty
+        return None, dict(answer_counts)
+
+
 # --- Sidebar Configuration ---
 with st.sidebar:
     st.header("⚙️ Core Settings")
@@ -389,47 +500,64 @@ with st.sidebar:
     with st.expander("🧠 Model Configuration", expanded=True):
         model_name = st.text_input(
             "Hugging Face Model ID or Path",
-            "ayjays132/NeuroReasoner-1-NR-1",
-            help="Enter the model ID from huggingface.co or a local path."
+            "ayjays132/NeuroReasoner-1-NR-1", # Default model
+            help="Enter the model ID from huggingface.co or a local path. Changing this requires reloading."
         )
 
         # --- Dynamic Model Type Detection ---
-        detected_type = "Unknown (Enter Model ID)"
-        # Options match the strings used in the loading function
-        model_type_options = ["Auto", "Causal", "Seq2Seq"]
-        default_model_type_index = model_type_options.index("Auto")
-
         # Attempt to load config to detect type without caching (lightweight check)
+        # This provides immediate feedback on the likely model type.
+        detected_type_display = "Unknown (Enter Model ID)"
+        initial_config = None
         try:
             if model_name and model_name.strip(): # Only attempt if input is not empty
-                initial_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, low_cpu_mem_usage=True)
-                is_encoder_decoder_initial = getattr(initial_config, "is_encoder_decoder", False)
-                detected_type = "Seq2Seq" if is_encoder_decoder_initial else "Causal"
+                with st.spinner("Detecting model type..."): # Small spinner for detection
+                    # Use from_pretrained without loading weights (low_cpu_mem_usage=True helps)
+                    initial_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, low_cpu_mem_usage=True)
+                    is_encoder_decoder_initial = getattr(initial_config, "is_encoder_decoder", False)
+                    detected_type_display = "Seq2Seq" if is_encoder_decoder_initial else "Causal"
+                    # Store the actual detected type for use in the selectbox default index
+                    actual_detected_type_for_index = "Seq2Seq" if is_encoder_decoder_initial else "Causal"
             else:
-                 detected_type = "Unknown (Enter Model ID)"
+                 detected_type_display = "Unknown (Enter Model ID)"
+                 actual_detected_type_for_index = "Auto" # Default to Auto index if no model name
+
         except Exception:
-            detected_type = "Unknown (Config Load Error)" # Indicate config load itself failed
+            detected_type_display = "Unknown (Config Load Error)" # Indicate config load itself failed
+            actual_detected_type_for_index = "Auto" # Default to Auto index if config load fails
+
+
+        # Options match the strings used in the loading function
+        model_type_options = ["Auto", "Causal", "Seq2Seq"]
+        # Set the default index based on the detected type, falling back to "Auto"
+        try:
+             default_model_type_index = model_type_options.index(actual_detected_type_for_index) if actual_detected_type_for_index in model_type_options else 0 # Default to Auto (index 0)
+        except ValueError:
+             default_model_type_index = 0 # Should not happen if logic is correct, but safety fallback
+
 
         forced_model_type = st.selectbox(
             "Architecture Type",
             model_type_options,
-            index=default_model_type_index,
-            help=f"Detected: {detected_type}. 'Auto' uses the detected type. Select manually if detection is incorrect or overridden."
+            index=default_model_type_index, # Use the detected type as the default selection
+            help=f"Detected: {detected_type_display}. 'Auto' uses the detected type. Select manually if detection is incorrect or overridden."
         )
 
         # --- Device Selection ---
+        # List available devices, prioritizing CUDA if available
         available_devices = ["cpu"]
         if torch.cuda.is_available():
-             available_devices.insert(0, "cuda") # Put cuda first if available
+            available_devices.insert(0, "cuda") # Put cuda first if available
 
         device = st.selectbox(
             "Device",
             available_devices,
+            index=(0 if "cuda" in available_devices else 0), # Default to cuda if available, else cpu
             help="Select the hardware device for computation (GPU recommended)."
         )
 
         st.markdown("""
-             <small>💡 Changing model settings requires reloading the model.</small>
+            <small>💡 Changing model settings requires reloading the model.</small>
         """, unsafe_allow_html=True)
 
 
@@ -439,74 +567,71 @@ with st.sidebar:
     st.markdown("Define how the AI generates reasoning steps and answers.")
 
     with st.expander("Basic Parameters", expanded=True):
-        # Finalized 'Number of Reasoning Chains' parameter
+        # Number of Reasoning Chains slider
         num_chains = st.slider(
             "Number of Reasoning Chains",
             min_value=1,
-            max_value=15, # Kept the higher max for more robustness
-            value=5,      # Kept the default of 5
-            help="How many independent reasoning chains to generate for analyzing the problem. More chains can improve Self-Consistency but take longer."
+            max_value=15, # Allow generating up to 15 chains
+            value=5,      # Default to 5 chains for a good balance
+            help="How many independent reasoning chains to generate for analyzing the problem. More chains can improve Self-Consistency but take longer and use more memory."
         )
 
-        # Finalized 'No-repeat Ngram Size' parameter
-        no_repeat_ngram_size = st.slider( # Using the standard name for GenerationConfig
-            "No-repeat Ngram Size",
-            min_value=0,
-            max_value=10,
-            value=3,
-            help="Avoids generating repeating sequences of N tokens. Set to 0 to disable."
-        )
-
-        # Self-Consistency checkbox remains
-        self_consistency = st.checkbox(
+        # Self-Consistency checkbox
+        self_consistency_enabled_gui = st.checkbox(
             "Enable Self-Consistency Voting",
-            value=True,
-            help="When enabled, the system generates multiple chains and identifies the most common final answer as the consensus. Requires 'Number of Reasoning Chains' > 1."
+            value=True, # Default to enabled
+            help="When enabled, the system generates multiple chains and identifies the most common final answer as the consensus via majority voting. Requires 'Number of Reasoning Chains' > 1."
         )
 
         # Conditional warning if Self-Consistency is on but num_chains is 1
-        if self_consistency and num_chains <= 1:
-             st.warning("Self-Consistency is most effective with 2 or more chains.") # Slightly rephrased warning
+        if self_consistency_enabled_gui and num_chains <= 1:
+             st.warning("Self-Consistency voting is most effective with 2 or more chains.")
 
 
     # Advanced Parameters
     with st.expander("🧪 Advanced Sampling Parameters"):
+        # Using standard parameter names from Hugging Face GenerationConfig
         max_new_tokens = st.slider(
-            "Max Tokens per Chain",
-            50, 2048, 768,
-            help="Maximum number of new tokens to generate for *each* individual reasoning chain. Adjust based on complexity expected."
+            "Max New Tokens per Chain",
+            50, 2048, 768, # Min, Max, Default
+            help="Maximum number of new tokens to generate for *each* individual reasoning chain. Adjust based on expected reasoning complexity and answer length."
         )
         temperature = st.slider(
             "Temperature",
-            0.0, 2.0, 0.8,
-            help="Controls the randomness of sampling. 0.0 is deterministic (greedy). Higher values increase diversity."
+            0.0, 2.0, 0.8, # Min, Max, Default
+            step=0.05, # Allow finer control
+            help="Controls the randomness of sampling. 0.0 is deterministic (greedy). Higher values increase diversity in reasoning paths."
         )
         top_k = st.slider(
             "Top-k",
-            0, 100, 50,
+            0, 200, 50, # Min, Max, Default (increased max k)
             help="Filter to consider only the top_k most likely tokens at each step (0 disables). Used with sampling."
         )
         top_p = st.slider(
             "Top-p (Nucleus Sampling)",
-            0.0, 1.0, 0.95,
+            0.0, 1.0, 0.95, # Min, Max, Default
+            step=0.01, # Allow finer control
             help="Filter to consider tokens with cumulative probability below top_p (0.0 disables). Used with sampling."
         )
+        repetition_penalty = st.slider(
+            "Repetition Penalty",
+            1.0, 2.0, 1.1, # Min, Max, Default
+            step=0.05, # Allow finer control
+            help="Penalizes repeated tokens or sequences. Higher values reduce repetition in the output."
+        )
+        no_repeat_ngram_size = st.slider(
+             "No-repeat Ngram Size",
+             0, 10, 0, # Min, Max, Default (changed default to 0, often less needed with repetition penalty)
+             help="Avoids repeating sequences of N tokens. Set to 0 to disable. Can help prevent loops in reasoning."
+        )
         do_sample = st.checkbox(
             "Enable Sampling",
-            value=True,
-            help="If checked, uses probabilistic sampling (controlled by Temperature, Top-k, Top-p). If unchecked, uses greedy decoding."
+            value=True, # Default to enabled
+            help="If checked, uses probabilistic sampling (controlled by Temperature, Top-k, Top-p). If unchecked, uses greedy decoding (deterministic)."
         )
         if not do_sample:
              st.info("Sampling disabled. Temperature, Top-k, and Top-p will be ignored.")
 
-        no_repeat_ngram_size = st.slider(
-            "No-repeat Ngram Size",
-            0, 10, 3,
-            help="Avoids repeating sequences of N tokens. Set to 0 to disable."
-        )
-        # Optional: Add a seed for reproducibility if desired
-        # generation_seed = st.number_input("Generation Seed (Optional)", value=-1, help="Set a positive integer for reproducible generation.")
-
 
     st.markdown("---") # Visual separator
 
@@ -528,15 +653,15 @@ with input_container:
     with prompt_col:
         prompt = st.text_area(
             "📝 Enter your query or problem:",
-            height=150,
+            height=180, # Increased height for better input experience
             placeholder="Example: If a train travels at 60 mph and a car at 40 mph, starting at the same time from cities 300 miles apart, how long until they meet? Think step-by-step.",
-            key="user_prompt" # Added key for stability
+            key="user_prompt" # Unique key for the widget
         )
 
     with button_col:
-        # Add some vertical space to align the button nicely
-        st.markdown("<div style='height: 3.5rem;'></div>", unsafe_allow_html=True)
-        run_button = st.button("✨ Generate Reasoning", use_container_width=True, key="generate_button") # Added key
+        # Add some vertical space to align the button nicely with the text area
+        st.markdown("<div style='height: 3.25rem;'></div>", unsafe_allow_html=True) # Adjusted height
+        run_button = st.button("✨ Generate Reasoning", use_container_width=True, key="generate_button") # Unique key
 
 # Container for status updates and results
 results_container = st.container()
@@ -546,223 +671,227 @@ results_container = st.container()
 if run_button:
     if not prompt or not prompt.strip():
         results_container.warning("Please enter a prompt to begin generation.")
-        st.stop() # Stop execution until prompt is entered
-
-    # --- Prepare for Generation ---
-    # Load model and tokenizer (handles caching internally with st.cache_resource
-    # via safe_load_model_with_status which also reports status)
-    # This happens only when the button is clicked and parameters might have changed
-    model, tokenizer, loaded_model_type = safe_load_model_with_status(model_name, device, forced_model_type)
-
-    if model is None or tokenizer is None:
-        # Error was already shown by safe_load_model_with_status
-        st.error("Model or tokenizer failed to load. Please check settings and traceback above.")
-        st.stop() # Stop if loading failed
-
-
-    # --- Configure Generation ---
-    # Use a status box for ongoing generation process
-    with results_container:
-        st.markdown("---") # Separator before results
-        generation_status = st.status("Preparing generation config...", expanded=True)
-        update_telemetry() # Update telemetry while status is active
-
-
-    try:
-        # Build GenerationConfig based on sidebar parameters
-        # num_return_sequences should match num_chains for the wrapper to process them
-        gen_cfg = GenerationConfig(
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p,
-            do_sample=do_sample,
-            num_return_sequences=num_chains, # <--- Corrected: Use num_chains directly
-            no_repeat_ngram_size=no_repeat_ngram_size,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.pad_token_id,
-            return_dict_in_generate=True,
-            output_scores=False,
-            output_attentions=False,
-            output_hidden_states=False,
-            use_cache=True,
-        )
-        generation_status.write(f"Generation parameters set: {gen_cfg.to_dict()}")
-        update_telemetry()
-
-        cfg = GenerationConfig(
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p,
-            do_sample=True,
-            num_return_sequences=num_chains,
-            no_repeat_ngram_size=no_repeat_ngram_size,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.pad_token_id
-        )
-    except Exception as e:
-        generation_status.error(f"❌ Failed to create GenerationConfig: {e}")
-        st.exception(e)
-        st.stop()
-
-    # --- Instantiate Wrapper ---
-    try:
-        generation_status.write("Initializing Chain-of-Thought wrapper...")
-        # Pass the configured generation config to the wrapper
-        # The wrapper should internally use num_return_sequences from gen_cfg
-        cot_wrapper = ChainOfThoughtWrapper(
-            model=model,
-            tokenizer=tokenizer,
-            generation_config=cfg,
-            device=device,
-            self_consistency=self_consistency,
-            consistency_rounds=(num_chains if self_consistency else 1)
-        )
-        generation_status.write("Wrapper initialized.")
-        update_telemetry()
-
-    except Exception as e:
-        generation_status.error(f"❌ Failed to initialize CoT wrapper: {e}")
-        st.exception(e)
-        st.stop()
-
-    # --- Tokenize Input ---
-    try:
-        generation_status.write("Tokenizing input prompt...")
-        # Use model_max_length or a reasonable cap for input length
-        max_input_length = tokenizer.model_max_length
-        if max_input_length is None or max_input_length > 4096: # Cap input length if tokenizer reports None or very large
-             max_input_length = 4096
-             if tokenizer.model_max_length is None:
-                  generation_status.warning(f"Tokenizer has no model_max_length, capping input to {max_input_length}.")
-
-
-        enc = tokenizer(
-            prompt,
-            return_tensors='pt',
-            padding='longest', # Pad to the longest sequence in the batch (batch size is 1 here)
-            truncation=True,
-            max_length=max_input_length, # Use a proper max length for the input
-        ).to(device)
-        generation_status.write(f"Input token length: {enc['input_ids'].shape[1]}")
-        update_telemetry()
-
-    except Exception as e:
-        generation_status.error(f"❌ Tokenization failed: {e}")
-        st.exception(e)
-        st.stop()
-
-    # --- Generate ---
-    generation_status.update(label=f"⏳ Generating {num_chains} reasoning chains...", state="running")
-    start_time = time.time()
+        # No need to st.stop() here, warning is sufficient
+    else:
+        # --- Prepare for Generation ---
+        # Load model and tokenizer (handles caching internally via safe_load_model_with_status)
+        # This happens only when the button is clicked and parameters might have changed
+        model, tokenizer, loaded_model_type = safe_load_model_with_status(model_name, device, forced_model_type)
+
+        if model is None or tokenizer is None:
+            # Error was already shown by safe_load_model_with_status
+            results_container.error("Model or tokenizer failed to load. Please check settings and traceback above.")
+            # No need to st.stop() here, error message is displayed
+
+        else: # Model and tokenizer loaded successfully
+            # --- Configure GenerationConfig ---
+            # Build the GenerationConfig object from sidebar parameters
+            # This config will be passed to the wrapper's __init__ as the base config
+            # and to the wrapper's generate() call as overrides.
+            # The wrapper's generate method will ultimately use these settings.
+            try:
+                # Base config for the wrapper's __init__ - defines default behavior
+                base_gen_config = GenerationConfig(
+                    max_new_tokens=max_new_tokens, # Use sidebar value
+                    temperature=temperature,       # Use sidebar value
+                    top_k=top_k,                   # Use sidebar value
+                    top_p=top_p,                   # Use sidebar value
+                    do_sample=do_sample,           # Use sidebar value
+                    repetition_penalty=repetition_penalty, # Use sidebar value
+                    no_repeat_ngram_size=no_repeat_ngram_size, # Use sidebar value
+                    eos_token_id=tokenizer.eos_token_id, # Always pass eos_token_id from tokenizer
+                    pad_token_id=tokenizer.pad_token_id, # Always pass pad_token_id from tokenizer
+                    # Other parameters like num_beams, etc., could be added here if exposed in sidebar
+                )
+
+
+            except Exception as e:
+                results_container.error(f"❌ Failed to create GenerationConfig from parameters: {e}")
+                st.exception(e)
+                st.stop() # Stop if config creation fails
+
+            # --- Instantiate Wrapper ---
+            # Use a status box for ongoing generation process
+            with results_container:
+                st.markdown("## ⏳ Generation Progress") # Use a clear header for the status section
+                generation_status = st.status("Initializing Chain-of-Thought wrapper...", expanded=True)
+                update_telemetry() # Update telemetry while status is active
 
-    try:
-        # Call the wrapper's generate method
-        # It should handle the loop for multiple chains and self-consistency internally
-        outputs = cot_wrapper.generate(
-            input_ids=enc['input_ids'],
-            attention_mask=enc['attention_mask'],
-            # Pass any other necessary arguments to your wrapper's generate method
-        )
-        # Expected `outputs` dict structure: {'full_texts': [...], 'reasoning_steps': [...], 'final_answers': [...], 'consensus_answer': '...'}
-        # The wrapper should handle extracting steps/answers if needed.
+                try:
+                    # Pass the loaded model, tokenizer, device, and relevant settings.
+                    # The wrapper uses the generation_config passed here as its base defaults.
+                    # Self-consistency settings from GUI are passed to wrapper's init attributes.
+                    cot_wrapper = ChainOfThoughtWrapper(
+                        model=model,
+                        tokenizer=tokenizer,
+                        # Pass the configured generation params as the base config for the wrapper
+                        generation_config=base_gen_config,
+                        device=device,
+                        # Pass GUI self-consistency settings
+                        # The wrapper uses self_consistency_enabled to decide if it *should* generate >1 chains
+                        # when num_return_sequences is not explicitly passed to generate().
+                        # We *are* explicitly passing num_return_sequences to generate(),
+                        # so these init flags primarily inform internal wrapper behavior/logging.
+                        self_consistency_enabled=self_consistency_enabled_gui,
+                        consistency_rounds=num_chains # Inform the wrapper about intended rounds
+                        # Pass other wrapper-specific init args if needed (e.g., custom tags)
+                        # final_answer_tag="Final Answer:" # Example if different from default
+                    )
+                    generation_status.write("Wrapper initialized.")
+                    update_telemetry()
+
+                except Exception as e:
+                    generation_status.error(f"❌ Failed to initialize CoT wrapper: {e}")
+                    st.exception(e)
+                    # Clean up resources in case of failure
+                    if torch.cuda.is_available(): torch.cuda.empty_cache()
+                    gc.collect()
+                    st.stop() # Stop if wrapper initialization fails
+
+                # --- Generate ---
+                # Call the wrapper's generate method.
+                # Pass the original input text.
+                # Explicitly pass num_return_sequences to request the desired number of chains.
+                generation_status.update(label=f"⏳ Generating {num_chains} reasoning chain(s)...", state="running")
+                start_time = time.time()
 
-    except Exception as e:
-        generation_status.error(f"❌ Generation failed: {e}")
-        st.exception(e)
-        # Clean up resources after potential OOM or other errors
-        if torch.cuda.is_available(): torch.cuda.empty_cache()
-        gc.collect() # Python garbage collection
-        st.stop()
-
-    elapsed_time = time.time() - start_time
-    generation_status.update(label=f"✨ Generation complete in {elapsed_time:.2f}s", state="complete")
-    update_telemetry() # Final telemetry update after successful generation
-
-    # --- Display Results ---
-    with results_container:
-        st.markdown("## 📚 Reasoning Output")
-
-        # Display Self-Consistency Consensus first if enabled and results are available
-        if self_consistency and outputs and 'consensus_answer' in outputs and outputs.get('final_answers'):
-            consensus = outputs.get('consensus_answer')
-            answers = outputs.get('final_answers', [])
-
-            st.markdown('<div class="consensus-answer">', unsafe_allow_html=True)
-            st.write("💡 **Consensus Answer (Self-Consistency):**")
-            st.write(consensus if consensus else "[Could not determine consensus]")
-            st.markdown('</div>', unsafe_allow_html=True)
-
-            if answers and len(answers) > 1: # Only show distribution if more than one answer was found
-                st.markdown("###### Answer Distribution:")
-                answer_counts = Counter(answers)
-                # Display sorted distribution
-                for ans, count in answer_counts.most_common():
-                     st.write(f"- '{ans}' ({count} {'vote' if count == 1 else 'votes'})")
-            st.markdown("---") # Separator
-
-
-        # Display individual chains
-        full_texts = outputs.get('full_texts', [])
-        reasoning_steps = outputs.get('reasoning_steps', [])
-        final_answers = outputs.get('final_answers', [])
-
-        if not full_texts:
-            st.warning("No reasoning chains were generated.")
-        else:
-            st.markdown(f"### Individual Chains ({len(full_texts)} generated)")
-            # Iterate and display each chain in an expander
-            # Ensure lists are iterable, even if empty
-            full_texts = full_texts if isinstance(full_texts, list) else []
-            reasoning_steps = reasoning_steps if isinstance(reasoning_steps, list) else []
-            final_answers = final_answers if isinstance(final_answers, list) else []
-
-            # Pad lists to the same length in case the wrapper returned inconsistent outputs
-            max_len_outputs = max(len(full_texts), len(reasoning_steps), len(final_answers))
-            full_texts.extend(["[N/A - Generation Failed for this chain]"] * (max_len_outputs - len(full_texts)))
-            reasoning_steps.extend([[]] * (max_len_outputs - len(reasoning_steps)))
-            final_answers.extend(["[N/A]"] * (max_len_outputs - len(final_answers)))
-
-
-            for idx, (text, steps, ans) in enumerate(zip(full_texts, reasoning_steps, final_answers), 1):
-                # Use try-except just in case a single chain output is malformed
                 try:
-                    # Expander for each chain, starting collapsed
-                    with st.expander(f"Chain {idx}", expanded=False):
-                        # Use custom class for styling the label
-                        st.markdown('<div class="output-label">Full Generated Text:</div>', unsafe_allow_html=True)
-                        # Use custom class for styling the text area background
-                        st.text_area(f"chain_text_area_{idx}", text, height=250, label_visibility="collapsed", help="The complete generated output for this chain.")
-
-                        if steps and isinstance(steps, list):
-                            st.markdown('<div class="output-label">Reasoning Steps:</div>', unsafe_allow_html=True)
-                            # Display steps as a list
-                            if steps:
-                                for i, step in enumerate(steps, 1):
-                                     if isinstance(step, str) and step.strip():
-                                          st.write(f"**Step {i}:** {step.strip()}")
-                                     elif not isinstance(step, str):
-                                          st.warning(f"Step {i} has invalid format.")
-                            else:
-                                 st.info("No specific steps were extracted for this chain.")
-
-
-                        st.markdown('<div class="output-label">Final Answer:</div>', unsafe_allow_html=True)
-                        st.write(f"**{ans if ans else '[No answer extracted]'}**")
-
-                        # Optional: Add a separator between chain sections
-                        st.markdown("---", help="End of Chain details.")
-
-                except Exception as chain_e:
-                    st.error(f"Error displaying Chain {idx}: {chain_e}")
-                    st.exception(chain_e)
-
-
-        st.markdown("---") # Final separator
-        st.info("Generation process concluded. Review the chains above.")
-
-    # Clean up GPU memory after generation is complete and results are displayed
-    if torch.cuda.is_available():
-         torch.cuda.empty_cache()
-    gc.collect() # Python garbage collection
+                    # The wrapper's generate takes input_text and optional config/num_return_sequences overrides.
+                    # We rely on the wrapper's internal config logic and pass the desired number of sequences.
+                    outputs = cot_wrapper.generate(
+                        input_text=prompt,
+                        # Optional: Pass a GenerationConfig override for this specific call if needed, e.g.:
+                        # generation_config=GenerationConfig(temperature=temperature + 0.1),
+                        # Pass the requested number of chains directly to the generate method:
+                        num_return_sequences=num_chains,
+                    )
+                    # Expected `outputs` dict structure: {'full_texts': [...], 'reasoning_steps': [...], 'final_answers': [...], 'generation_scores': [...]}
+
+                except Exception as e:
+                    generation_status.error(f"❌ Generation failed: {e}")
+                    st.exception(e) # Display the full exception traceback
+                    # Clean up resources after potential OOM or other errors
+                    if torch.cuda.is_available(): torch.cuda.empty_cache()
+                    gc.collect() # Python garbage collection
+                    # No need to st.stop() here, error message is displayed
+                    outputs = None # Ensure outputs is None if generation fails
+
+                elapsed_time = time.time() - start_time
+
+                # Update final status based on success or failure
+                if outputs is not None:
+                    generation_status.update(label=f"✨ Generation complete in {elapsed_time:.2f}s", state="complete")
+                else:
+                     generation_status.update(label=f"❌ Generation failed after {elapsed_time:.2f}s", state="error")
+
+                update_telemetry() # Final telemetry update after generation attempt
+
+
+            # --- Process and Display Results ---
+            if outputs is not None: # Only display if generation was attempted and returned results
+                with results_container:
+                    st.markdown("## 📚 Reasoning Output")
+
+                    # Display Self-Consistency Consensus first if enabled and results are available
+                    # Implement the voting logic here in the GUI using the wrapper's output
+                    final_answers_list = outputs.get('final_answers', [])
+
+                    if self_consistency_enabled_gui and final_answers_list:
+                        # Perform the actual voting
+                        consensus_answer, answer_distribution_dict = perform_self_consistency_voting(final_answers_list)
+                        # Convert dict to Counter for sorting by count in display
+                        answer_distribution = Counter(answer_distribution_dict)
+
+
+                        st.markdown('<div class="consensus-answer">', unsafe_allow_html=True)
+                        st.write("💡 **Consensus Answer (Self-Consistency):**")
+                        if consensus_answer:
+                             st.write(f'<p>{consensus_answer}</p>', unsafe_allow_html=True)
+                             # Optional: Add a note about the confidence/number of votes for the winner
+                             # We need the count of the winning answer from the distribution
+                             winner_count = answer_distribution.get(normalize_answer(consensus_answer), 0) # Get count for the winning normalized answer
+                             st.write(f"*(Based on {winner_count} {'vote' if winner_count == 1 else 'votes'} out of {len(final_answers_list)} chains)*")
+                        else:
+                             st.write("<p>[Could not determine consensus - no valid answers found]</p>", unsafe_allow_html=True)
+                             st.write(f"*(Examined {len(final_answers_list)} chains)*")
+
+                        st.markdown('</div>', unsafe_allow_html=True)
+
+                        # Display answer distribution if there's more than one unique answer after normalization
+                        if len(answer_distribution) > 1:
+                            st.markdown("###### Answer Distribution:")
+                            # Display sorted distribution by vote count
+                            for ans, count in answer_distribution.most_common():
+                                # Display the normalized answer and its count
+                                st.write(f"- '{ans}' ({count} {'vote' if count == 1 else 'votes'})")
+                        elif len(answer_distribution) == 1 and consensus_answer:
+                             st.info(f"All {len(final_answers_list)} valid chains agreed on the normalized answer: '{consensus_answer}'.")
+                        else:
+                             st.warning(f"No valid answers ({len(answer_distribution)} unique normalized answers) were found to determine distribution from {len(final_answers_list)} chains.")
+
+
+                        st.markdown("---") # Separator after consensus section
+
+
+                    # Display individual chains
+                    full_texts = outputs.get('full_texts', [])
+                    reasoning_steps_list = outputs.get('reasoning_steps', [])
+                    final_answers_list_raw = outputs.get('final_answers', []) # Keep raw answers for display
+
+                    if not full_texts:
+                        st.warning("No reasoning chains were generated or parsed successfully.")
+                    else:
+                        st.markdown(f"### Individual Chains ({len(full_texts)} generated)")
+                        # Iterate and display each chain in an expander
+                        # Ensure lists are iterable and have consistent length, padding with placeholders if necessary
+                        max_len_outputs = len(full_texts)
+                        # Ensure reasoning_steps_list and final_answers_list_raw match the length of full_texts
+                        reasoning_steps_list = (reasoning_steps_list if isinstance(reasoning_steps_list, list) else []) + [[]] * (max_len_outputs - len(reasoning_steps_list))
+                        final_answers_list_raw = (final_answers_list_raw if isinstance(final_answers_list_raw, list) else []) + ["[N/A - Parsing Failed]"] * (max_len_outputs - len(final_answers_list_raw))
+
+
+                        for idx, (text, steps, ans_raw) in enumerate(zip(full_texts, reasoning_steps_list, final_answers_list_raw), 1):
+                            # Use try-except for displaying each chain just in case of unexpected data format
+                            try:
+                                # Expander for each chain, starting collapsed
+                                with st.expander(f"Chain {idx}", expanded=False):
+                                    # Display full generated text
+                                    st.markdown('<div class="output-label">Full Generated Text (Cleaned):</div>', unsafe_allow_html=True)
+                                    st.text_area(f"chain_text_area_{idx}", text if isinstance(text, str) else "[Invalid Text Data]", height=250, label_visibility="collapsed", help="The complete generated output for this chain after cleaning artifacts.", key=f"chain_text_{idx}") # Added key
+
+                                    # Display parsed reasoning steps
+                                    st.markdown('<div class="output-label">Reasoning Steps Parsed:</div>', unsafe_allow_html=True)
+                                    if steps and isinstance(steps, list) and len(steps) > 0:
+                                        # Display steps as a list with strong emphasis on step number
+                                        for i, step in enumerate(steps, 1):
+                                            if isinstance(step, str) and step.strip():
+                                                 st.markdown(f"**Step {i}:** {step.strip()}")
+                                            elif not isinstance(step, str):
+                                                 st.warning(f"Step {i} has invalid format in Chain {idx}.")
+                                    elif isinstance(steps, list) and len(steps) == 0:
+                                         st.info("No specific steps were extracted for this chain.")
+                                    else:
+                                         st.warning(f"Reasoning steps data is invalid or missing for Chain {idx}.")
+
+
+                                    # Display parsed final answer (raw)
+                                    st.markdown('<div class="output-label">Final Answer Parsed:</div>', unsafe_allow_html=True)
+                                    display_answer = ans_raw if isinstance(ans_raw, str) and ans_raw.strip() else "[No answer extracted]"
+                                    st.write(f"**{display_answer}**")
+
+                                    # Optional: Add a separator between chain sections within the expander if desired
+                                    # st.markdown("---")
+
+                            except Exception as chain_e:
+                                st.error(f"Error displaying content for Chain {idx}: {chain_e}")
+                                st.exception(chain_e)
+
+                        # Final separator after all chains
+                        st.markdown("---")
+                        st.info(f"Displayed details for {len(full_texts)} generated chains.")
+
+            # Clean up GPU memory after generation and display are complete
+            if torch.cuda.is_available():
+                 torch.cuda.empty_cache()
+            gc.collect() # Python garbage collection
+            # st.write("GPU memory cache cleared and garbage collected.") # Optional status message

chain_of_thought_wrapper.py

@@ -3,47 +3,60 @@
 import re
 import torch
 import logging
-from transformers import PreTrainedModel, AutoTokenizer, GenerationConfig, GenerationMixin
+from transformers import (
+    PreTrainedModel,
+    AutoTokenizer,
+    GenerationConfig,
+    GenerationMixin,
+    AutoModelForCausalLM # Needed for example usage
+)
 from transformers.utils import is_accelerate_available, is_bitsandbytes_available
 from typing import Optional, List, Tuple, Dict, Union, Any
 import gc # Import garbage collector for cleanup
-import time
+import time # Import time for potential timing/logging (unused in final code, but good practice)
+from collections import Counter # Needed for example voting
 
 # --- Logging Setup ---
-# Configure logging for the module
-logging.basicConfig(level=logging.INFO) # Default logging level
+# Configure logging for the module. This helps in debugging and understanding wrapper behavior.
+# Set level to DEBUG temporarily to see the detailed logs added below
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
-# Prevent duplicate handlers if imported multiple times
+# Ensure logger doesn't add handlers multiple times if the script is imported repeatedly
 if not logger.handlers:
     handler = logging.StreamHandler()
     formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
     handler.setFormatter(formatter)
     logger.addHandler(handler)
-    logger.propagate = False # Prevent logs from going to root logger multiple times
-
+    # Avoid propagation to the root logger, preventing duplicate messages
+    logger.propagate = False
 
 # --- Default Configuration Values ---
-# These defaults provide sensible starting points for the wrapper's behavior.
-DEFAULT_MAX_LENGTH = 1024          # Default maximum length of the generated output sequence.
-DEFAULT_REASONING_LIMIT = 10       # Limit on the number of steps to extract during parsing (currently unused in parse logic, but good to keep as a concept).
-DEFAULT_CONSISTENCY_ROUNDS = 3     # Default number of chains to generate for self-consistency (used in __init__, passed via GUI num_chains).
-DEFAULT_COMPLEXITY_KEYWORDS = ["explain", "step by step", "plan", "analyze", "reasoning", "logic"] # Keywords to potentially trigger CoT (currently unused, CoT is always on).
-DEFAULT_FINAL_ANSWER_TAG = "Final_Answer:" # The specific tag expected before the final answer.
+# These defaults provide sensible starting points for the wrapper's behavior,
+# based on common practices and the audit recommendations.
+DEFAULT_MAX_LENGTH = 2048 # Increased default max length to accommodate longer CoT
+DEFAULT_REASONING_LIMIT = 15 # A conceptual limit for extracted steps (not strictly enforced by parsing logic)
+DEFAULT_CONSISTENCY_ROUNDS = 5 # Default number of chains for self-consistency, increased based on typical research
+DEFAULT_COMPLEXITY_KEYWORDS = ["explain", "step by step", "plan", "analyze", "reasoning", "logic"] # Keywords (currently unused as CoT is always on)
+DEFAULT_FINAL_ANSWER_TAG = "Final_Answer:" # Explicit tag to signal the final answer
 
 # --- Regex Pattern for Parsing Steps ---
 # This pattern is used to identify and extract individual reasoning steps from
-# the generated text. It's designed to be flexible, capturing:
-# - "Step N:"
-# - "Step N."
-# - "Step N-"
-# - "N:"
-# - "N."
-# - "N-"
-# Where N is one or more digits, case-insensitive for "Step".
+# the generated text. It's designed to be flexible, capturing common step formats
+# like "Step N:", "N.", etc., case-insensitive for "Step".
+# Captures the text *after* the step marker.
 DEFAULT_STEP_PATTERN = re.compile(
-    r"^(?:Step\s*\d+[:.)-]|\d+[:.)-])\s*(.*)", re.IGNORECASE
+    r"^(?:Step\s*\d+[:.)-]\s*|\d+[:.)-]\s*)(.*)", re.IGNORECASE
 )
 
+# --- Common Artifact Cleanup Regex ---
+# Regex patterns to remove common problematic tokens or structures models sometimes emit,
+# which are not part of the desired reasoning or answer. Based on audit suggestion.
+ARTIFACT_PATTERNS = [
+    re.compile(r"<init>.*?</init>", re.DOTALL),       # Example: DeepSeek R1 init tags
+    re.compile(r"<final_output>.*?</final_output>", re.DOTALL), # Example: DeepSeek R1 final output tags
+    # re.compile(r"\{.*?\}", re.DOTALL), # Removing all {} might be too aggressive, removed based on re-evaluation.
+    # Add other specific artifact patterns here as needed for observed model outputs
+]
 
 class ChainOfThoughtWrapper:
     """
@@ -53,15 +66,20 @@ class ChainOfThoughtWrapper:
     template into the prompt. It handles model generation and parses the
     output to extract reasoning steps and a final answer. It is designed
     to generate multiple sequences for potential Self-Consistency voting
-    (voting logic is expected to be handled by the calling application,
-    like the Streamlit GUI).
+    (voting logic is intended for the calling application, e.g., a GUI).
+
+    It incorporates enhancements based on a detailed audit, focusing on
+    prompting, decoding, parsing robustness, cross-model compatibility,
+    reliability mitigation, and efficiency, while adhering to the "always-on CoT"
+    principle.
 
     Key Features:
-    - Forces CoT via prompt injection.
-    - Parses structured reasoning steps and final answer from output.
-    - Supports generating multiple chains for Self-Consistency analysis.
-    - Compatible with Hugging Face PreTrainedModels or objects implementing `.generate()`.
-    - Handles device placement and merges GenerationConfig.
+    - Forces CoT via a structured, adaptive prompt template.
+    - Parses structured reasoning steps and uses robust logic to find the final answer.
+    - Supports generating multiple chains for Self-Consistency analysis via GenerationConfig.
+    - Handles common cross-model compatibility issues (e.g., pad tokens, device placement).
+    - Merges user-provided GenerationConfig with sensible defaults.
+    - Includes basic cleanup for common model output artifacts.
     """
 
     def __init__(
@@ -71,143 +89,179 @@ class ChainOfThoughtWrapper:
         generation_config: Optional[GenerationConfig] = None,
         device: Optional[str] = None,
         max_length: int = DEFAULT_MAX_LENGTH,
-        reasoning_steps_limit: int = DEFAULT_REASONING_LIMIT, # Parameter included as per provided code
-        self_consistency: bool = False, # Parameter included as per provided code (__init__ attribute)
-        consistency_rounds: int = DEFAULT_CONSISTENCY_ROUNDS, # Parameter included as per provided code (__init__ attribute)
-        complexity_keywords: Optional[List[str]] = None, # Parameter included as per provided code
+        reasoning_steps_limit: int = DEFAULT_REASONING_LIMIT,
+        self_consistency_enabled: bool = False, # Control if multiple chains are generated
+        consistency_rounds: int = DEFAULT_CONSISTENCY_ROUNDS,
+        complexity_keywords: Optional[List[str]] = None, # Currently unused as CoT is always on
         final_answer_tag: str = DEFAULT_FINAL_ANSWER_TAG,
-        # self_consistency_enabled: bool = False # Removed this based on user's 'keep as is' and gui interaction
+        # Optional prompt customization for advanced users
+        cot_instruction: str = "Let's analyze this problem logically, breaking it down step by step to reach the precise final answer.",
+        reasoning_header: str = "Reasoning Process:",
+        step_prefix: str = "Step ", # e.g., "Step 1: " - model will ideally continue this
+        # Optional reliability controls (simple, prompt-based)
+        emphasize_factual: bool = True,
+        allow_uncertainty_phrase: Optional[str] = "If information is insufficient or you are unsure, state that clearly.",
+        # Optional parsing flexibility
+        strip_artifact_patterns: List[re.Pattern] = ARTIFACT_PATTERNS,
     ):
         """
-        Initializes the ChainOfThoughtWrapper.
+        Initializes the ChainOfThoughtWrapper with enhanced configurations.
 
         Args:
             model (Union[PreTrainedModel, GenerationMixin, Any]): The language model.
-                                                                   Must have a `.generate()` method.
+                                                                   Must have a .generate() method.
             tokenizer (AutoTokenizer): The corresponding tokenizer.
             generation_config (Optional[GenerationConfig]): A default generation configuration.
-                                                            Values here can be overridden by `generate()` call.
+                                                            Values here can be overridden by generate() call.
             device (Optional[str]): The device to load the model onto ('cpu' or 'cuda').
                                     Defaults to 'cuda' if available, otherwise 'cpu'.
             max_length (int): The maximum total length of the input + generated sequence.
-            reasoning_steps_limit (int): Conceptual limit for parsed steps (currently not enforced in _parse).
-            self_consistency (bool): Flag indicating if self-consistency is intended (Informs `consistency_rounds` attribute).
-            consistency_rounds (int): The number of chains to generate if self-consistency is active (Informs `consistency_rounds` attribute).
-                                      The actual number generated is controlled by `num_return_sequences` in `generate()` or `generation_config`.
-            complexity_keywords (Optional[List[str]]): List of keywords to potentially trigger CoT (currently unused).
+                              This should be large enough for the prompt, reasoning, and answer.
+            reasoning_steps_limit (int): Conceptual limit for parsed steps. Not strictly enforced by current parsing.
+            self_consistency_enabled (bool): If True, enable multi-chain generation for self-consistency.
+            consistency_rounds (int): The number of chains to generate if `self_consistency_enabled` is True.
+                                      Actual number of sequences is controlled by `num_return_sequences`
+                                      in the final `GenerationConfig`.
+            complexity_keywords (Optional[List[str]]): List of keywords (unused with always-on CoT).
             final_answer_tag (str): The specific string marker expected before the final answer.
+            cot_instruction (str): The core instruction phrase for CoT.
+            reasoning_header (str): The header text before the reasoning steps.
+            step_prefix (str): The prefix for the first step.
+            emphasize_factual (bool): If True, add prompt text emphasizing factual reasoning.
+            allow_uncertainty_phrase (Optional[str]): If provided, add a phrase prompting model to state uncertainty.
+            strip_artifact_patterns (List[re.Pattern]): List of regex patterns to remove from model output before parsing.
         """
-        # Determine and set the device
+        # --- Device Handling ---
+        # Determine and set the device. Log the chosen device.
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-        logger.info("Initializing wrapper on device: %s", self.device)
+        logger.info("Initializing ChainOfThoughtWrapper on device: %s", self.device)
 
-        # Move the model to the specified device
+        # --- Model and Tokenizer Loading and Configuration ---
+        # Move the model to the specified device and set to evaluation mode.
+        # Includes error handling for device transfer.
         try:
             self.model = model.to(self.device)
-            self.model.eval() # Set model to evaluation mode for consistent behavior
+            self.model.eval() # Set model to evaluation mode (disables dropout, etc.)
             logger.info("Model moved to %s and set to eval mode.", self.device)
         except Exception as e:
             logger.error("Failed to move model to device %s: %s", self.device, e)
-            raise # Re-raise the exception after logging
+            raise # Re-raise the exception if device transfer fails
 
         self.tokenizer = tokenizer
 
-        # Set core parameters
+        # Attempt to find the underlying Hugging Face model instance and its config.
+        # This helps reliably access attributes like `config.vocab_size`, `resize_token_embeddings`, etc.
+        self._hf_model_instance, self._hf_config = self._find_hf_model_and_config(self.model)
+
+        # Handle models/tokenizers without a defined pad_token_id.
+        # This is crucial for batch generation (like `num_return_sequences`).
+        # If the tokenizer doesn't have a pad_token, try to use the eos_token.
+        # If neither exists, add a special token and resize embeddings.
+        # The wrapper's `resize_token_embeddings` method is called here if a new token is added.
+        if self.tokenizer.pad_token_id is None:
+             if self.tokenizer.eos_token_id is not None:
+                  self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+                  logger.warning("Tokenizer pad_token_id is None, using eos_token_id (%s) as pad_token_id.", self.tokenizer.eos_token_id)
+             else:
+                 # Fallback: Add a new pad token if neither exists
+                 logger.warning("Tokenizer pad_token_id and eos_token_id are both None. Adding a [PAD] token.")
+                 self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+                 self.tokenizer.pad_token_id = self.tokenizer.convert_tokens_to_ids('[PAD]')
+                 logger.info("Added new [PAD] token with ID %s.", self.tokenizer.pad_token_id)
+                 # Resize model embeddings if we added a new token AND we found a base HF model instance
+                 if self._hf_model_instance:
+                     self.resize_token_embeddings(len(self.tokenizer)) # Call the instance method
+                     logger.info("Resized model embeddings to accommodate new PAD token.")
+                 else:
+                     logger.warning("Could not resize model embeddings after adding PAD token; underlying HF model instance not found.")
+                     logger.warning("Ensure the model can handle a larger vocabulary if batching is used.")
+
+        # --- Configuration Attributes ---
         self.max_length = max_length
         self.reasoning_steps_limit = reasoning_steps_limit
-        self.self_consistency = self_consistency # Attribute stored, actual generation count controlled elsewhere
-        self.consistency_rounds = max(1, consistency_rounds) if self_consistency else 1 # Attribute stored
-        self.complexity_keywords = complexity_keywords or list(DEFAULT_COMPLEXITY_KEYWORDS) # Ensure it's a mutable list
+        # The actual number of sequences to generate is controlled by `num_return_sequences` in the final `GenerationConfig`.
+        # We store `consistency_rounds` to potentially inform this value.
+        self.self_consistency_enabled = self_consistency_enabled
+        self.consistency_rounds = max(1, consistency_rounds) if self_consistency_enabled else 1
+
+        # --- Prompt Template Components ---
+        self.complexity_keywords = complexity_keywords or list(DEFAULT_COMPLEXITY_KEYWORDS) # Store keywords (currently unused for logic)
         self.final_answer_tag = final_answer_tag
-        # Compile regex pattern for final answer extraction
+        self._cot_instruction = cot_instruction # Customizable CoT instruction
+        self._reasoning_header = reasoning_header # Customizable reasoning header
+        self._step_prefix = step_prefix # Customizable step prefix (e.g., "Step ")
+
+        # --- Reliability/Hallucination Mitigation Prompt Components ---
+        self._emphasize_factual = emphasize_factual
+        self._allow_uncertainty_phrase = allow_uncertainty_phrase
+
+        # --- Parsing Attributes and Compiled Regex ---
+        # Compile regex pattern for final answer extraction based on the specified tag.
+        # re.escape handles potential special characters in the tag. re.DOTALL matches newline.
         self.final_answer_pattern = re.compile(
             re.escape(final_answer_tag) + r"\s*(.*)", re.IGNORECASE | re.DOTALL
         )
-        logger.debug("Final answer pattern compiled: %s", self.final_answer_pattern.pattern)
-        logger.debug("Step pattern: %s", DEFAULT_STEP_PATTERN.pattern)
-
-        # Attempt to find the underlying Hugging Face model and its config
-        # This is useful for accessing standard attributes like eos_token_id, etc.
-        self._hf_model, self._hf_config = self._find_hf_model_and_config(self.model)
-
-        # Fallback to tokenizer settings if HF config isn't found
-        if self._hf_config is None:
-            logger.warning("Underlying HF model config not found. Relying on tokenizer for eos/pad tokens and vocab size.")
-            # Create a pseudo-config with essential tokenizer info
-            class PseudoConfig:
-                def __init__(self, tok):
-                    self.eos_token_id = tok.eos_token_id
-                    # Use eos_token_id as pad_token_id if pad_token_id is None (common for GPT-like models)
-                    self.pad_token_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
-                    # Fallback if both are None (less common but possible)
-                    if self.pad_token_id is None:
-                        logger.warning("Tokenizer pad_token_id and eos_token_id are both None. Generation might be unstable without padding.")
-                        # Assign a arbitrary value or handle externally if this happens in practice
-                        # For now, keep it None, generation might fail or behave unexpectedly
-                        pass # Keep pad_token_id as None
-
-                    self.vocab_size = len(tok) # Vocabulary size from tokenizer
-
-                def __getattr__(self, name):
-                     # Allow accessing other attributes, returning None if not found
-                     # This prevents errors if generation_config tries to read something unexpected
-                    logger.debug("Accessing undefined attribute '%s' on PseudoConfig. Returning None.", name)
-                    return None
-
-            self._hf_config = PseudoConfig(self.tokenizer)
-            logger.debug("Created PseudoConfig: eos_token_id=%s, pad_token_id=%s, vocab_size=%s",
-                         self._hf_config.eos_token_id, self._hf_config.pad_token_id, self._hf_config.vocab_size)
-        else:
-            logger.info("Found underlying HF model config.")
-            logger.debug("HF Config: eos_token_id=%s, pad_token_id=%s, vocab_size=%s",
-                         getattr(self._hf_config, 'eos_token_id', None),
-                         getattr(self._hf_config, 'pad_token_id', None),
-                         getattr(self._hf_config, 'vocab_size', None))
+        self._step_pattern = DEFAULT_STEP_PATTERN # Use the default compiled step pattern
+        self._artifact_patterns = strip_artifact_patterns # Patterns for cleaning model output
 
+        logger.debug("Final answer pattern compiled: %s", self.final_answer_pattern.pattern)
+        logger.debug("Step pattern: %s", self._step_pattern.pattern)
 
-        # --- Setup Generation Config ---
-        # Start with a base config, either provided or a default one
+        # --- Base Generation Config Setup ---
+        # Create or copy the base GenerationConfig. This config holds the default
+        # generation parameters that will be used unless overridden during a generate() call.
+        # Use .from_dict(.to_dict()) for a clean copy if a config was provided.
         if generation_config:
-            # Use from_dict and to_dict for safe merging/copying of GenerationConfig
-            self.generation_config = GenerationConfig.from_dict(generation_config.to_dict())
-            logger.info("Initialized with provided GenerationConfig.")
+            self.base_generation_config = GenerationConfig.from_dict(generation_config.to_dict())
+            logger.info("Initialized with provided base GenerationConfig.")
         else:
-            # Create a default GenerationConfig using info from HF config or tokenizer fallback
-            self.generation_config = GenerationConfig(
-                eos_token_id=self._hf_config.eos_token_id,
-                pad_token_id=self._hf_config.pad_token_id,
-                max_length=self.max_length, # Set max_length from wrapper param
-                # Add other common defaults if not provided
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.95,
-                top_k=50,
-                num_return_sequences=1, # Default to 1 sequence
-                no_repeat_ngram_size=0, # Default to no ngram repetition prevention
+            # Create a default GenerationConfig if none was provided.
+            # Incorporate parameters known to work well for CoT based on audit (temp, top_p, top_k).
+            # Ensure pad_token_id and eos_token_id are set from the tokenizer (or the fallback).
+            self.base_generation_config = GenerationConfig(
+                eos_token_id=self.tokenizer.eos_token_id,
+                pad_token_id=self.tokenizer.pad_token_id,
+                max_length=self.max_length, # Max total length
+                do_sample=True,             # Always sample for diversity (essential for multi-chain)
+                temperature=0.7,            # Balanced randomness
+                top_p=0.95,                 # Nucleus sampling
+                top_k=50,                   # Top-k sampling cutoff
+                num_return_sequences=1,     # Default to 1 sequence (will be overridden by generate call if self-consistency is on)
+                # Add a mild repetition penalty, useful for longer CoT
+                repetition_penalty=1.1,     # Discourage immediate repetition
+                no_repeat_ngram_size=0,     # Default to no n-gram repetition prevention
             )
-            logger.info("Initialized with default GenerationConfig.")
-
-        # Ensure the underlying HF model (if found) is set to return dict outputs from generate
-        # This is necessary for accessing scores, hidden states etc. if needed, and for consistency.
-        # Use a check as some custom models might not have this attribute on their config.
-        if hasattr(self._hf_config, 'return_dict_in_generate'):
-            try:
-                setattr(self._hf_config, 'return_dict_in_generate', True)
-                logger.debug("Set _hf_config.return_dict_in_generate = True.")
-            except Exception as e:
-                logger.warning("Failed to set return_dict_in_generate on _hf_config: %s", e)
+            logger.info("Initialized with default base GenerationConfig.")
+
+        # Ensure the base config uses the determined pad_token_id
+        # This might be redundant if tokenizer already has it, but ensures consistency
+        self.base_generation_config.pad_token_id = self.tokenizer.pad_token_id
+        logger.debug("Base GenerationConfig pad_token_id set to %s.", self.base_generation_config.pad_token_id)
+
+        # Check if the underlying HF model (if found) supports returning scores, useful for CISC.
+        # We set this on the model's config if possible, as `generate` reads from there.
+        if self._hf_model_instance and hasattr(self._hf_model_instance.config, 'return_dict_in_generate'):
+             try:
+                 # Set these attributes directly on the model's config object
+                 self._hf_model_instance.config.return_dict_in_generate = True
+                 self._hf_model_instance.config.output_scores = True # Also request scores
+                 logger.debug("Set underlying HF model config to return dict in generate and output scores.")
+             except Exception as e:
+                 logger.warning("Failed to set return_dict_in_generate/output_scores on HF model config: %s", e)
         else:
-             logger.debug("_hf_config does not have return_dict_in_generate attribute.")
+             logger.debug("Underlying HF model instance or config does not support setting return_dict_in_generate/output_scores.")
 
 
-        logger.info("ChainOfThoughtWrapper initialization complete on device: %s", self.device)
-        logger.debug("Initial GenerationConfig: %s", self.generation_config.to_dict())
+        logger.info("ChainOfThoughtWrapper initialization complete.")
+        logger.debug("Final Base GenerationConfig: %s", self.base_generation_config.to_dict())
 
 
     def _find_hf_model_and_config(self, obj: Any) -> Tuple[Optional[PreTrainedModel], Optional[Any]]:
         """
         Recursively searches for an underlying Hugging Face PreTrainedModel
-        and its configuration within a potentially wrapped object.
+        and its configuration within a potentially wrapped or custom object.
+        This helps in accessing standard HF attributes like `config` or
+        methods like `resize_token_embeddings`.
 
         Args:
             obj (Any): The object to inspect (could be the model itself or a wrapper).
@@ -216,14 +270,21 @@ class ChainOfThoughtWrapper:
             Tuple[Optional[PreTrainedModel], Optional[Any]]: The found HF model instance and its config.
                                                               Returns (None, None) if not found.
         """
+        # Add a check to prevent infinite recursion
+        if getattr(obj, '_searching_hf_model', False):
+             logger.debug("Preventing infinite recursion in _find_hf_model_and_config for object type: %s", type(obj))
+             return None, None
+        setattr(obj, '_searching_hf_model', True)
+
         logger.debug("Searching for HF model in object of type: %s", type(obj))
         # If the object is directly a PreTrainedModel and has a config
-        if isinstance(obj, PreTrainedModel) and hasattr(obj, 'config'):
+        if isinstance(obj, PreTrainedModel):
             logger.debug("Found HF PreTrainedModel directly.")
-            return obj, obj.config
+            setattr(obj, '_searching_hf_model', False) # Reset flag
+            return obj, getattr(obj, 'config', None) # Return config if it exists
 
         # Check common attribute names where the base model might be stored
-        potential_attrs = ('model', 'base_model', 'transformer', 'hf_model')
+        potential_attrs = ('model', 'base_model', 'transformer', '_original_model', 'module') # Added 'module'
         for attr_name in potential_attrs:
             m = getattr(obj, attr_name, None)
             if m is not None:
@@ -231,22 +292,25 @@ class ChainOfThoughtWrapper:
                 # Recursively search within the attribute
                 found_model, found_config = self._find_hf_model_and_config(m)
                 if found_model or found_config:
+                    setattr(obj, '_searching_hf_model', False) # Reset flag before returning
                     return found_model, found_config
 
-        # If no PreTrainedModel found, check if the object itself has a 'config' attribute
+        # If no PreTrainedModel found through attributes, check if the object itself has a 'config' attribute
         if hasattr(obj, 'config'):
-            logger.debug("Found config attribute on object, but no PreTrainedModel.")
-            return None, obj.config
+             logger.debug("Found config attribute on object, but no PreTrainedModel instance.")
+             setattr(obj, '_searching_hf_model', False) # Reset flag
+             return None, obj.config # Return the config found
 
-        logger.debug("No HF PreTrainedModel or config found.")
+        logger.debug("No underlying HF PreTrainedModel instance or config found.")
+        setattr(obj, '_searching_hf_model', False) # Reset flag
         return None, None
 
 
     def _inject_cot(self, prompt: str) -> str:
         """
-        Injects the prescriptive Chain-of-Thought template into the user's prompt.
-
-        This method defines the expected format the model should follow for reasoning.
+        Injects the structured Chain-of-Thought template into the user's prompt.
+        This template guides the model's response format.
+        Incorporates reliability prompts based on settings.
 
         Args:
             prompt (str): The original user prompt.
@@ -254,178 +318,239 @@ class ChainOfThoughtWrapper:
         Returns:
             str: The prompt with the CoT template appended.
         """
-        # The template strongly guides the model to produce step-by-step reasoning
-        # followed by a specific tag for the final answer.
-        cot_prompt = (
-            f"{prompt.strip()}\n\n" # Use strip() to clean user prompt
-            "Let's analyze this problem logically, breaking it down step by step to reach the precise final answer.\n\n" # Enhanced instruction
-            "Reasoning Process:\n\n" # Clearer heading for steps
-            "Step 1: " # Start the first step explicitly
-            # More steps are not needed here, the model learns to continue the pattern
-        )
-        logger.debug("Injected CoT template. Full prompt starts with: %s...", cot_prompt[:100].replace('\n', '\\n'))
-        return cot_prompt
+        # Start with the cleaned original prompt
+        injected_prompt = f"{prompt.strip()}\n\n"
 
+        # Add the core CoT instruction phrase
+        injected_prompt += self._cot_instruction + "\n"
 
-    @torch.no_grad()
+        # Add reliability-focused instructions if enabled
+        if self._emphasize_factual:
+             injected_prompt += "Think through the problem step-by-step using only factual information and logical deduction. Do not assume any facts that are not given.\n"
+        if self._allow_uncertainty_phrase:
+             injected_prompt += self._allow_uncertainty_phrase + "\n"
+
+        # Add the structured template for reasoning steps and final answer tag
+        injected_prompt += f"\n{self._reasoning_header}\n\n"
+        injected_prompt += f"{self._step_prefix}1: " # Explicitly start the first step to guide format consistency
+
+        logger.debug("Injected CoT template. Full prompt starts with: %s...", injected_prompt[:200].replace('\n', '\\n'))
+        return injected_prompt
+
+
+    @torch.no_grad() # Disable gradient calculation during generation for efficiency
     def generate(
         self,
-        input_ids: torch.LongTensor,
-        attention_mask: Optional[torch.LongTensor] = None,
-        generation_config: Optional[GenerationConfig] = None,
-        num_return_sequences: int = 1, # This argument controls how many sequences are generated
-        **kwargs: Any # Allows passing arbitrary generation parameters
+        input_text: str,
+        generation_config: Optional[GenerationConfig] = None, # Optional override config for this call
+        num_return_sequences: Optional[int] = None, # Explicitly request N sequences
     ) -> Dict[str, Any]:
         """
-        Generates text using the wrapped model, enforcing Chain-of-Thought.
-
-        This method prepares the input by injecting the CoT template, calls the
-        underlying model's generate method, and then parses the raw outputs
-        to extract structured reasoning steps and final answers.
+        Generates text using the wrapped model with Chain-of-Thought injection.
+        Handles tokenization, prompt injection, generation, and parsing.
+        Efficiently generates multiple sequences using `num_return_sequences`.
 
         Args:
-            input_ids (torch.LongTensor): Tokenized input prompt (batch size 1 expected).
-                                          Shape [1, sequence_length].
-            attention_mask (Optional[torch.LongTensor]): Attention mask for the input.
-                                                        Shape [1, sequence_length].
-            generation_config (Optional[GenerationConfig]): Specific generation config
-                                                           for this call. Overrides defaults.
-            num_return_sequences (int): The number of independent sequences to generate.
-                                        This is crucial for Self-Consistency.
-                                        Comes from the GUI's 'num_chains'.
-            **kwargs (Any): Additional keyword arguments passed to the model's `generate` method.
+            input_text (str): The user's input text/question.
+            generation_config (Optional[GenerationConfig]): Additional generation parameters
+                                                            to override the base config for this call.
+            num_return_sequences (Optional[int]): Number of independent sequences (chains) to generate.
+                                                  If None, uses the value from the merged generation config
+                                                  (defaulting to 1 or `consistency_rounds` if enabled).
 
         Returns:
-            Dict[str, Any]: A dictionary containing:
-                            - 'sequences' (torch.LongTensor): The raw generated token sequences.
-                            - 'full_texts' (List[str]): The complete decoded text for each sequence.
-                            - 'reasoning_steps' (List[List[str]]): List of parsed reasoning steps for each sequence.
-                            - 'final_answers' (List[str]): List of parsed final answers for each sequence.
-                            - 'consensus_answer' (Optional[str]): The consensus answer if self-consistency is active and possible (Handled by calling code).
+            Dict[str, Any]: A dictionary containing the generation results:
+                            - 'sequences': The raw generated token IDs (list of tensors).
+                            - 'full_texts': List of raw, cleaned text outputs (after stripping prompt/artifacts) for each chain.
+                            - 'reasoning_steps': List of lists of extracted reasoning steps for each chain.
+                            - 'final_answers': List of extracted final answer strings for each chain.
+                            - 'generation_scores': Scores if requested and available (for CISC externally).
         """
-        # Ensure input is on the correct device
-        input_ids = input_ids.to(self.device)
-        if attention_mask is not None:
-            attention_mask = attention_mask.to(self.device)
-
-        # Decode the original prompt text for CoT injection
-        # Assume batch size is 1 for the input prompt tensor [1, sequence_length]
-        if input_ids.size(0) != 1:
-             logger.warning("Batch size > 1 detected for input_ids (%d). CoT injection assumes batch size 1. Using the first item.", input_ids.size(0))
-        prompt_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
-
-        # --- Inject CoT Template ---
-        # This is the core step that forces the model into a reasoning mode.
-        cot_prompt = self._inject_cot(prompt_text)
-        logger.debug("Injected CoT prompt. Encoding...")
-
-        # --- Prepare Generation Configuration ---
-        # Merge the wrapper's default config with the call-specific config and kwargs.
-        # The num_return_sequences from the function argument takes precedence here.
-        cfg = GenerationConfig.from_dict(self.generation_config.to_dict()) # Start with wrapper's default
-        if generation_config:
-             cfg.update(**generation_config.to_dict()) # Update with call-specific config
-             logger.debug("Updated GenerationConfig with call-specific config.")
-
-        # Explicitly set num_return_sequences from the function argument
-        cfg.num_return_sequences = num_return_sequences
-        logger.info("Generating %d sequence(s).", cfg.num_return_sequences)
-
-        # Update with any remaining keyword arguments passed to generate()
-        for k, v in kwargs.items():
-             if hasattr(cfg, k):
-                  setattr(cfg, k, v)
-                  logger.debug("Updating GenerationConfig kwarg: %s=%s", k, v)
-             else:
-                 # Allow passing arbitrary kwargs to model.generate if the underlying method supports them
-                 # These won't be part of the GenerationConfig object itself unless it's a supported param.
-                 # However, the model's generate method might accept extra args.
-                 # Log a warning if it's not a standard GenerationConfig parameter.
-                 if k not in GenerationConfig().__dict__: # Check if it's NOT a standard param
-                     logger.debug("Passing non-standard kwarg '%s' to model.generate.", k)
-                 # We pass all kwargs to model.generate below anyway.
-
-        logger.debug("Final GenerationConfig for call: %s", cfg.to_dict())
-
-        # --- Encode the CoT Prompt ---
-        # Max length for input should be total max_length minus max_new_tokens
-        # to leave space for the generation.
-        # Ensure padding and truncation are handled.
+        logger.info("Received generate call with input text starting: '%s...'", input_text[:100])
+
+        # 1) Inject the CoT prompt into the original input text
+        cot_prompt_text = self._inject_cot(input_text)
+
+        # 2) Tokenize the full CoT prompt
+        # Ensure padding is handled correctly. Use return_tensors="pt" for PyTorch tensors.
+        # truncation=True ensures the input fits within max_length.
+        # max_length applies to the input sequence here.
         try:
-            enc = self.tokenizer(
-                cot_prompt,
-                return_tensors='pt',
-                padding='longest', # Pad to the longest sequence in the batch (always 1 here)
-                truncation=True,   # Crucially, truncate if the prompt is too long
-                max_length=self.max_length - cfg.max_new_tokens # Leave room for generation
+            encoded_input = self.tokenizer(
+                cot_prompt_text,
+                return_tensors="pt",
+                padding="longest", # Pad to the longest sequence in the batch (only 1 here, but good practice)
+                truncation=True,
+                max_length=self.max_length, # Truncate if the prompt itself is too long
             ).to(self.device)
-            logger.debug("Encoded CoT prompt. Input shape: %s", enc['input_ids'].shape)
-
+            logger.debug("Input text tokenized. Input IDs shape: %s, on device: %s", encoded_input['input_ids'].shape, encoded_input['input_ids'].device)
         except Exception as e:
-            logger.error("Failed to encode CoT prompt: %s", e)
-            raise # Re-raise the exception after logging
+            logger.error("Failed to tokenize input text: %s", e)
+            # Attempt cleanup before raising
+            if torch.cuda.is_available(): torch.cuda.empty_cache()
+            gc.collect()
+            raise # Re-raise tokenization error
+
+
+        # 3) Build the final GenerationConfig for this specific call
+        # Start with the base config, then merge any provided overrides.
+        # Use .from_dict(.to_dict()) for safe merging.
+        cfg = GenerationConfig.from_dict(self.base_generation_config.to_dict())
+
+        if generation_config is not None:
+             logger.debug("Merging provided generation_config overrides...")
+             cfg.update(**generation_config.to_dict())
+             logger.debug("Merged user-provided GenerationConfig.")
+
+        # Explicitly set num_return_sequences for this call based on the argument.
+        # This overrides any num_return_sequences set in the base config or the provided override config.
+        if num_return_sequences is not None:
+            cfg.num_return_sequences = num_return_sequences
+            logger.debug("Using num_return_sequences from function argument: %s", cfg.num_return_sequences)
+        elif self.self_consistency_enabled:
+             # Fallback: If num_return_sequences argument is None, use consistency_rounds if self_consistency is enabled
+             cfg.num_return_sequences = self.consistency_rounds
+             logger.debug("num_return_sequences argument is None, using consistency_rounds (%s) because self_consistency is enabled.", cfg.num_return_sequences)
+        else:
+             # Fallback: If num_return_sequences argument is None and self_consistency is disabled, default to 1
+             cfg.num_return_sequences = 1
+             logger.debug("num_return_sequences argument is None and self_consistency disabled, defaulting to 1.")
+
+
+        # Ensure max_length in the config respects the wrapper's max_length setting
+        # max_length in generate() config is the *total* length (input + new tokens)
+        # max_new_tokens is the number of *new* tokens generated
+        # Prefer max_new_tokens if set, otherwise calculate from max_length
+        input_length = encoded_input['input_ids'].shape[1]
+        if cfg.max_new_tokens is None:
+             # If max_new_tokens is NOT set, ensure the total length does not exceed the wrapper's max_length
+             if cfg.max_length is not None:
+                 # Only adjust cfg.max_length if it's set in the base/override config
+                 cfg.max_length = min(self.max_length, cfg.max_length)
+             else:
+                 # If neither max_new_tokens nor max_length were set in base/override, use wrapper's max_length
+                 cfg.max_length = self.max_length
+             logger.debug("max_new_tokens not set in config. Using total max_length: %s (Input length: %s)", cfg.max_length, input_length)
+        else:
+            # If max_new_tokens IS set, the total length will be input_length + max_new_tokens
+            # We should check if this effective total length exceeds the wrapper's overall max_length
+            effective_total_length = input_length + cfg.max_new_tokens
+            if effective_total_length > self.max_length:
+                 logger.warning("Effective total length (input %d + new %d = %d) exceeds wrapper max_length (%d). Adjusting max_new_tokens.",
+                                input_length, cfg.max_new_tokens, effective_total_length, self.max_length)
+                 # Adjust max_new_tokens down to respect the wrapper's limit
+                 cfg.max_new_tokens = max(0, self.max_length - input_length)
+                 logger.warning("Adjusted max_new_tokens to %d.", cfg.max_new_tokens)
+
+        # Ensure pad_token_id and eos_token_id are correctly set in the final config
+        # Use tokenizer's IDs as the source of truth
+        cfg.pad_token_id = self.tokenizer.pad_token_id
+        cfg.eos_token_id = self.tokenizer.eos_token_id
+
+        logger.debug("Final GenerationConfig for this call after resolving overrides and num_return_sequences: %s", cfg.to_dict())
+
+        # --- Debugging: Inspect inputs immediately before generation ---
+        # ADDED LOGGING HERE TO DIAGNOSE CUDA ERROR
+        logger.debug("-" * 30 + " Inputs to model.generate " + "-" * 30)
+        logger.debug("  Input Text Snippet: '%s...'", input_text[:100])
+        logger.debug("  CoT Prompt Text Snippet: '%s...'", cot_prompt_text[:200].replace('\n', '\\n'))
+        logger.debug("  Input IDs shape: %s, dtype: %s, device: %s", encoded_input["input_ids"].shape, encoded_input["input_ids"].dtype, encoded_input["input_ids"].device)
+        if encoded_input.get("attention_mask", None) is not None:
+            logger.debug("  Attention Mask shape: %s, dtype: %s, device: %s", encoded_input["attention_mask"].shape, encoded_input["attention_mask"].dtype, encoded_input["attention_mask"].device)
+            # Log a snippet of the attention mask for inspection (only first batch item, first 20 tokens)
+            if encoded_input["attention_mask"].numel() > 0:
+                 logger.debug("  Attention Mask snippet (first 20): %s", encoded_input["attention_mask"][0, :20].tolist())
+            # Check if mask seems valid (contains only 0s and 1s) - might not catch all CUDA errors but helps debug
+            if not torch.all((encoded_input["attention_mask"] == 0) | (encoded_input["attention_mask"] == 1)):
+                 logger.error("!!! Attention mask contains values other than 0 or 1 !!!")
+        else:
+            logger.warning("!!! No attention mask provided to model.generate !!!")
+        logger.debug("  GenerationConfig.pad_token_id: %s", cfg.pad_token_id)
+        logger.debug("  GenerationConfig.eos_token_id: %s", cfg.eos_token_id)
+        logger.debug("  GenerationConfig.num_return_sequences: %s", cfg.num_return_sequences)
+        logger.debug("-" * 30 + " End Inputs to model.generate " + "-" * 30)
+        # --- End Debugging ---
 
 
-        # --- Generate Text ---
-        # Call the underlying model's generate method with the prepared input and config.
-        # torch.no_grad() context is already applied to the whole method.
+        # 4) Generate text using the model's generate method
+        # Pass input_ids and attention_mask. Pass the *final* GenerationConfig object.
         try:
-            logger.info("Calling model.generate()...")
-            start_time = time.time() # Measure generation time
-            out = self.model.generate(
-                input_ids=enc['input_ids'],
-                attention_mask=enc['attention_mask'],
-                generation_config=cfg,
-                **kwargs # Pass through any extra kwargs
+            generation_output = self.model.generate(
+                input_ids=encoded_input["input_ids"],
+                attention_mask=encoded_input.get("attention_mask", None),
+                generation_config=cfg, # Pass the fully configured GenerationConfig
+                # Request scores if supported by the model/config for potential CISC implementation externally
+                return_dict_in_generate=True, # Request dict output
+                output_scores=True,          # Request scores
             )
-            elapsed_time = time.time() - start_time
-            logger.info("model.generate() finished in %.2f seconds.", elapsed_time)
-            logger.debug("Raw output shape: %s", out.shape)
-
+            generated_sequences = generation_output.sequences
+            # If scores were requested and returned, they are available in generation_output.scores
+            # These can be used by the caller for CISC voting.
+            generation_scores = generation_output.scores if hasattr(generation_output, 'scores') else None
+            logger.info("Generation complete. Generated %d sequence(s).", len(generated_sequences))
+            if generation_scores:
+                 logger.debug("Generation scores available (%d scores tensors).", len(generation_scores))
 
         except Exception as e:
             logger.error("Model generation failed: %s", e)
-            # Attempt to clean up GPU memory in case of OOM or other errors
-            if torch.cuda.is_available():
-                 torch.cuda.empty_cache()
-            gc.collect() # Trigger Python garbage collection
-            raise # Re-raise the exception after logging
+            # Log the exception details
+            import traceback
+            logger.error(traceback.format_exc()) # Log full traceback
 
-        # --- Decode and Parse Outputs ---
-        # Decode the generated token sequences back into text.
-        logger.debug("Decoding and parsing outputs...")
-        decoded_outputs = self.tokenizer.batch_decode(out, skip_special_tokens=True)
-
-        # Process each decoded output to extract steps and final answer
-        parsed_results = [self._parse(text, cot_prompt) for text in decoded_outputs]
-
-        # Separate the parsed components into lists
-        all_steps = [r[0] for r in parsed_results]
-        all_finals = [r[1] for r in parsed_results]
-        all_full_texts = [r[2] for r in parsed_results] # The 'body' after removing template
-
-        logger.info("Generated and parsed %d sequence(s).", len(decoded_outputs))
-
-        # --- Return Results ---
-        # The calling code (e.g., the GUI) is responsible for implementing
-        # Self-Consistency voting based on the list of 'final_answers' provided here.
+             # Attempt cleanup even on failure - this *might* also trigger the CUDA error again,
+             # but it's the correct place to *try* to clean up GPU memory associated with the model.
+            if torch.cuda.is_available():
+                try:
+                    torch.cuda.empty_cache()
+                    logger.debug("Attempted torch.cuda.empty_cache() after generation failure.")
+                except Exception as cache_e:
+                    logger.error("Error during cuda empty_cache after generation failure: %s", cache_e)
+            gc.collect()
+            logger.debug("Attempted gc.collect() after generation failure.")
+
+            raise # Re-raise generation error
+
+
+        # 5) Decode and Parse the generated sequences
+        # Ensure generated_sequences is a list or tensor before decoding
+        if not isinstance(generated_sequences, (list, torch.Tensor)) or len(generated_sequences) == 0:
+             logger.warning("No sequences generated. Returning empty results.")
+             return {
+                "sequences": [],
+                "full_texts": [],
+                "reasoning_steps": [],
+                "final_answers": [],
+                "generation_scores": None,
+             }
+
+        decoded_outputs = self.tokenizer.batch_decode(generated_sequences, skip_special_tokens=True)
+        logger.debug("Batch decoding complete.")
+        parsed_results = [self._parse(text, cot_prompt_text) for text in decoded_outputs]
+        logger.debug("Parsing complete for %d sequences.", len(parsed_results))
+
+
+        # Unpack the parsed results
+        all_steps = [result[0] for result in parsed_results]
+        all_final_answers = [result[1] for result in parsed_results]
+        full_generated_bodies = [result[2] for result in parsed_results]
+
+        # 6) Construct and return the results dictionary
+        # The actual self-consistency voting logic is handled by the caller,
+        # but the wrapper provides the necessary outputs (multiple chains and parsed answers).
         return {
-            'sequences': out, # Return raw sequences in case they are needed
-            'full_texts': all_full_texts, # Text body after template removal
-            'reasoning_steps': all_steps,
-            'final_answers': all_finals,
-            # 'consensus_answer' is not computed here, it's done externally.
-            # Keeping the structure consistent with GUI expectation.
-            'consensus_answer': None # Placeholder, computed externally
+            "sequences": generated_sequences, # Raw sequences (token IDs)
+            "full_texts": full_generated_bodies, # Cleaned generated text bodies
+            "reasoning_steps": all_steps, # Parsed reasoning steps for each chain
+            "final_answers": all_final_answers, # Parsed final answer for each chain
+            "generation_scores": generation_scores, # Scores if requested and available (for CISC)
         }
 
 
     def _parse(self, text: str, cot_prompt: str) -> Tuple[List[str], str, str]:
         """
         Parses the generated text to extract reasoning steps and the final answer.
-
-        Applies regex patterns to find lines matching the step format and the
-        final answer tag. Includes cleanup for stray model artifacts.
+        This is a robust parsing function that handles different formats,
+        artifacts, and provides fallback logic for finding the answer.
 
         Args:
             text (str): The raw text output from the model for a single chain.
@@ -435,229 +560,368 @@ class ChainOfThoughtWrapper:
             Tuple[List[str], str, str]: A tuple containing:
                                         - A list of extracted reasoning step strings.
                                         - The extracted final answer string.
-                                        - The full body of the generated text (after removing the prompt).
+                                        - The full body of the generated text (after removing the prompt and artifacts).
         """
-        logger.debug("Parsing generated text...")
+        logger.debug("Starting parsing for a single generated text chunk...")
 
-        # Remove the exact injected prompt from the beginning of the text.
+        # 1) Remove the exact injected prompt from the beginning of the text.
         # This isolates the model's generated continuation.
         body = text
         if text.startswith(cot_prompt):
-            body = text[len(cot_prompt):].strip()
-            logger.debug("Removed CoT prompt (%d characters) from beginning.", len(cot_prompt))
+            body = text[len(cot_prompt):] # Remove the prefix
+            logger.debug("Removed exact CoT prompt (%d characters) from beginning.", len(cot_prompt))
+        else:
+            logger.warning("Generated text does not start with the injected CoT prompt. Attempting to parse entire text after initial whitespace strip.")
+            body = text.lstrip() # Just strip leading whitespace if template wasn't followed
+
+        # 2) Apply artifact cleanup patterns
+        logger.debug("Applying artifact cleanup patterns...")
+        original_body_len = len(body)
+        cleaned_body = body # Start with body after prompt removal
+        for pattern in self._artifact_patterns:
+            cleaned_body = pattern.sub("", cleaned_body)
+        if len(cleaned_body) < original_body_len:
+            logger.debug("Artifact cleanup removed %d characters.", original_body_len - len(cleaned_body))
         else:
-            logger.warning("Generated text does not start with the injected CoT prompt. Parsing entire text.")
-            body = text.strip() # Just strip whitespace if template wasn't followed
-
-        # --- Cleanup stray model artifacts ---
-        # Remove common problematic tags or partial JSON structures that models sometimes emit.
-        # This makes the raw output cleaner before step/answer extraction.
-        logger.debug("Cleaning stray artifacts...")
-        body = re.sub(r"<init>.*?</init>", "", body, flags=re.DOTALL)
-        body = re.sub(r"<final_output>.*?</final_output>", "", body, flags=re.DOTALL)
-        # Note: Removing all {} might be aggressive if model uses them naturally.
-        # Keeping it as per provided code, but be aware this could remove desired output.
-        # Consider making this optional or more specific if needed.
-        body = re.sub(r"\{.*?\}", "", body, flags=re.DOTALL)
-        logger.debug("Artifact cleanup complete.")
-
-
-        lines = [l.strip() for l in body.splitlines() if l.strip()] # Split into non-empty, stripped lines
+            logger.debug("No artifacts found matching patterns.")
+
+        # Ensure body is stripped after cleanup
+        cleaned_body = cleaned_body.strip()
+        body_lines = [l.strip() for l in cleaned_body.splitlines() if l.strip()] # Split into non-empty, stripped lines
+
         steps = [] # List to store extracted steps
         final_answer = "" # Variable to store the final answer
+        found_final_answer_tagged = False # Flag to track if the specific tag was found
 
-        # --- Extract Steps and Final Answer ---
+        # 3) Extract Steps and Final Answer (Primary Method: Tagged Answer)
         # Iterate through lines and apply regex patterns.
-        found_final_answer_line = False
-        for i, line in enumerate(lines):
+        # Prioritize finding the explicit final answer tag.
+        logger.debug("Attempting to extract steps and final answer using explicit tag '%s'...", self.final_answer_tag)
+        for i, line in enumerate(body_lines):
+            # Check for the explicit final answer tag pattern first
+            final_answer_match = self.final_answer_pattern.search(line)
+            if final_answer_match:
+                final_answer = final_answer_match.group(1).strip()
+                logger.debug("Extracted final answer using explicit tag: '%s'", final_answer[:100])
+                found_final_answer_tagged = True
+                # Once the tagged answer is found, we can stop processing lines for it
+                # We still iterate through ALL lines below to capture all steps BEFORE the tag.
+                # No break here because we need to collect steps that might appear after the tag was first encountered on a line.
+                # E.g., "Step 1: ... Final_Answer: X Step 2: ..." (unlikely but possible)
+                # The logic below ensures we capture steps *before* the final answer.
+
+
+        # Now, iterate through lines AGAIN to collect steps.
+        # This second pass ensures we collect steps even if the answer tag was found early.
+        # We stop collecting steps once we encounter the line that *contained* the final answer tag,
+        # or if we apply a step limit.
+        logger.debug("Collecting reasoning steps...")
+        for i, line in enumerate(body_lines):
+            # Stop collecting steps if we found the final answer tag on this line or a previous one
+            # And if we've reached or passed the line where the tag was found (if it was found)
+            # This requires knowing the index of the line where the tag was found.
+            # A simpler approach: just collect all lines matching step pattern UP TO the first line
+            # where the final answer tag was found.
+            final_answer_line_index = -1
+            for idx, l in enumerate(body_lines):
+                 if self.final_answer_pattern.search(l):
+                      final_answer_line_index = idx
+                      break # Found the first occurrence of the tag
+
+            if final_answer_line_index != -1 and i >= final_answer_line_index:
+                 logger.debug("Stopped collecting steps at line index %d because final answer tag was found on line %d.", i, final_answer_line_index)
+                 break # Stop collecting steps once we reach the line with the answer tag
+
             # Check for reasoning step pattern
-            step_match = DEFAULT_STEP_PATTERN.match(line)
+            step_match = self._step_pattern.match(line)
             if step_match:
-                # If a step is found, add the captured group (the text after the number/tag)
-                steps.append(step_match.group(1).strip())
-                logger.debug("Extracted step %d: '%s'", len(steps), steps[-1][:50])
-                # Stop adding steps if we've reached a defined limit (though limit isn't currently enforced after parsing)
-                # if len(steps) >= self.reasoning_steps_limit:
-                #    logger.debug("Reached reasoning steps limit (%d). Stopping step extraction.", self.reasoning_steps_limit)
-                #    # Continue iterating to potentially find the final answer after the limit
-                #    # break # DO NOT break if we still need to find the final answer tag after the limit
-
-
-            else:
-                # If it's not a step, check for the final answer tag
-                final_answer_match = self.final_answer_pattern.search(line)
-                if final_answer_match:
-                    # If the final answer tag is found, extract the text following it
-                    final_answer = final_answer_match.group(1).strip()
-                    logger.debug("Extracted final answer tagged: '%s'", final_answer[:50])
-                    found_final_answer_line = True
-                    # Once the final answer tag is found, we can stop processing lines for *this specific pattern*
-                    # However, the provided code breaks the loop entirely here.
-                    # Keeping the break to match the original logic.
-                    break # Stop processing lines after finding the tagged answer
-
-        # --- Fallback for Final Answer ---
-        # If the specific final answer tag was not found, assume the last non-step line
-        # is the intended final answer. This is a heuristic fallback.
-        if not found_final_answer_line:
-            logger.debug("Final answer tag not found. Applying fallback heuristic.")
-            # Find the last line that is not a step
+                step_text = step_match.group(1).strip()
+                if step_text: # Only add non-empty steps
+                    steps.append(step_text)
+                    # logger.debug("Extracted step: '%s'", steps[-1][:50]) # Too verbose usually
+                # Stop adding steps if we've reached a defined limit
+                if len(steps) >= self.reasoning_steps_limit:
+                   logger.debug("Reached reasoning steps limit (%d). Stopping step extraction.", self.reasoning_steps_limit)
+                   break # Stop collecting steps if limit is reached
+
+
+        # 4) Fallback for Final Answer (If Tag Still Not Found)
+        # If the explicit final answer tag was not found after both passes, apply fallback heuristics.
+        if not found_final_answer_tagged:
+            logger.debug("Explicit final answer tag not found. Applying fallback heuristics.")
+
+            # Fallback: Assume the last non-step line is the answer.
+            # Iterate backwards through the processed lines to find the last line that doesn't look like a step.
+            # Using the 'body_lines' list after cleanup and stripping.
             last_non_step_line = ""
-            for line in reversed(lines): # Iterate backwards
-                if line.strip() and not DEFAULT_STEP_PATTERN.match(line):
+            for line in reversed(body_lines): # Iterate backwards through non-empty, stripped lines
+                if line and not self._step_pattern.match(line):
                     last_non_step_line = line.strip()
-                    logger.debug("Fallback: Last non-step line found: '%s'", last_non_step_line[:50])
-                    break # Found the last non-step line
+                    logger.debug("Fallback: Identified last non-step line: '%s'", last_non_step_line[:100])
+                    break # Found the last non-step line, stop searching backwards
 
             if last_non_step_line:
                  # Check if the last non-step line *contains* the final answer tag,
-                 # even if it didn't *start* with it or was the last line processed.
-                 # This handles cases where the tag might be mid-line or in a different format.
+                 # even if it didn't *start* with it or wasn't the line where the tag was first found.
                  fa_match_fallback = self.final_answer_pattern.search(last_non_step_line)
                  if fa_match_fallback:
                       final_answer = fa_match_fallback.group(1).strip()
-                      logger.debug("Fallback found tagged answer in last non-step line: '%s'", final_answer[:50])
+                      logger.debug("Fallback found tagged answer in last non-step line: '%s'", final_answer[:100])
                  else:
-                    # If no tag in the last non-step line, just use the line itself
+                    # If no tag in the last non-step line, just use the line itself as the answer
                     final_answer = last_non_step_line
-                    logger.debug("Fallback using last non-step line as answer: '%s'", final_answer[:50])
+                    logger.debug("Fallback using last non-step line as answer: '%s'", final_answer[:100])
             else:
-                 # If no non-empty lines were found, the final answer is empty
+                 # If no non-empty or non-step lines were found, the final answer is empty
                  final_answer = ""
-                 logger.debug("No lines found in body. Final answer is empty.")
+                 logger.debug("Fallback: No non-empty or non-step lines found in body. Final answer is empty.")
+
+        # 5) Basic Post-Parsing Cleanup on Final Answer
+        # Remove any trailing punctuation from the final answer, unless it's part of specific symbols (like !?)
+        # This helps normalize answers for voting.
+        if final_answer:
+            # Remove common trailing characters like periods, commas, etc.
+            final_answer = re.sub(r'[.,;:]+$', '', final_answer).strip()
+            # Remove common leading "Answer: " or similar preambles if they weren't removed by tag matching
+            # This needs to be case-insensitive
+            final_answer = re.sub(r'^\s*(?:Answer|Result|Output|Final Answer)\s*[:\-]?\s*', '', final_answer, flags=re.IGNORECASE).strip()
+            logger.debug("Applied basic post-parsing cleanup to final answer: '%s'", final_answer[:100])
 
+        # Final check: Ensure steps list doesn't contain the final answer line or text
+        # This is a belt-and-suspenders approach as the logic above should prevent it,
+        # but safeguards against edge cases where the tag wasn't found but the line
+        # looked like a step *and* contained the answer.
+        if final_answer and steps:
+             # Remove any step that exactly matches the final answer after stripping
+             steps = [step for step in steps if step.strip() != final_answer.strip()]
+             # Also check if the final answer is contained *within* a step (less likely but possible)
+             steps = [step for step in steps if final_answer.strip() not in step.strip()]
 
-        logger.debug("Parsing complete. Steps found: %d, Final Answer: '%s'", len(steps), final_answer[:50])
 
-        return steps, final_answer, body # Return steps, final answer, and the cleaned body text
+        logger.info("Parsing complete. Steps found: %d, Final Answer: '%s'", len(steps), final_answer[:100])
+
+        # Return the extracted steps, the final answer, and the cleaned generated body text
+        return steps, final_answer, cleaned_body # Return steps, final answer, and the cleaned body text
 
 
     def resize_token_embeddings(self, new_size: int):
         """
-        Resizes the model's token embeddings, useful after adding new tokens
-        to the tokenizer (like a custom PAD token).
+        Resizes the model's token embeddings to match a new vocabulary size,
+        useful after adding new tokens (like a custom PAD token) to the tokenizer.
+        This operation is crucial if the tokenizer size changes and the model
+        is used for generation or training.
 
-        Only works if the underlying model object has a `resize_token_embeddings` method.
+        Only works if the underlying model object is a PreTrainedModel
+        or has a `resize_token_embeddings` method.
 
         Args:
             new_size (int): The new size of the vocabulary/embedding layer.
-                            Should match the size of the tokenizer's vocabulary.
+                            Should typically be `len(self.tokenizer)`.
         """
-        # Find the actual HF model if wrapped
-        hf_model_instance, _ = self._find_hf_model_and_config(self.model)
+        # Use the stored HF model instance found during initialization
+        hf_model_instance = self._hf_model_instance
 
-        if hasattr(hf_model_instance, 'resize_token_embeddings'):
+        if hf_model_instance and hasattr(hf_model_instance, 'resize_token_embeddings'):
             try:
                 old_size = hf_model_instance.get_input_embeddings().weight.size(0)
                 if new_size != old_size:
+                    logger.info("Attempting to resize model token embeddings from %d to %d.", old_size, new_size)
+                    # Ensure the model is on the correct device before resizing
+                    hf_model_instance.to(self.device)
                     hf_model_instance.resize_token_embeddings(new_size)
-                    logger.info("Resized model token embeddings from %d to %d.", old_size, new_size)
+                    logger.info("Successfully resized token embeddings.")
                     # Update model config's vocab size if available
                     if hasattr(hf_model_instance, 'config') and hasattr(hf_model_instance.config, 'vocab_size'):
-                        hf_model_instance.config.vocab_size = new_size
-                        logger.debug("Updated model config vocab_size to %d.", new_size)
+                         hf_model_instance.config.vocab_size = new_size
+                         logger.debug("Updated underlying model config vocab_size to %d.", new_size)
+                    # Attempt garbage collection after a potentially memory-intensive operation
+                    if torch.cuda.is_available(): torch.cuda.empty_cache()
+                    gc.collect()
                 else:
                     logger.info("Embedding size is already %d, no resizing needed.", new_size)
             except Exception as e:
                  logger.error("Failed to resize token embeddings: %s", e)
-                 # Attempt cleanup
+                 # Attempt cleanup even on failure
                  if torch.cuda.is_available(): torch.cuda.empty_cache()
                  gc.collect()
+                 # Note: Not re-raising here by default, as a failure might not be critical
+                 # depending on the user's intended use (e.g., if they don't use the new tokens for generation).
+                 # Could be re-raised if this is deemed a critical error.
         else:
-            logger.error("Cannot resize token embeddings: The underlying model object does not have a 'resize_token_embeddings' method.")
+            logger.warning("Cannot resize token embeddings: The underlying model object does not have a 'resize_token_embeddings' method or HF model instance not found.")
 
 
-# Example Usage (Illustrative - requires a real HF model and tokenizer)
+# Example Usage (Illustrative)
 if __name__ == "__main__":
     print("--- ChainOfThoughtWrapper Example Usage ---")
-    print("This block requires a Hugging Face model to run.")
-    print("Loading a small dummy model for demonstration...")
+    print("This block demonstrates loading a small HF model and using the wrapper.")
+    print("Setting logging level to DEBUG to see detailed wrapper logs.")
+    logger.setLevel(logging.DEBUG) # Set logger to DEBUG for example
 
     # You would replace this with your actual model loading logic
     try:
         # Use a tiny, fast model for a quick test
-        model_id = "hf-internal-testing/tiny-random-gpt2"
+        # NOTE: distilgpt2 might still hit CUDA errors with num_return_sequences > 1
+        # if there are underlying driver/CUDA/PyTorch compatibility issues or
+        # subtle model-specific padding bugs in HF transformers for this architecture.
+        # If this example still fails, try a different simple causal model like 'gpt2' or a small LLaMA variant.
+        model_id = "distilbert/distilgpt2" # A slightly larger but still fast GPT-2 variant
         device = "cuda" if torch.cuda.is_available() else "cpu"
 
         logger.info(f"Attempting to load model {model_id} on {device}...")
+
+        # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForCausalLM.from_pretrained(model_id)
 
-        # Ensure pad token is set for generation (common requirement)
+        # Ensure pad token is set for generation robustness (common requirement for GPT-like models)
+        # Handle this *before* loading the model if possible, or ensure embeddings are resized.
         if tokenizer.pad_token_id is None:
              if tokenizer.eos_token_id is not None:
                   tokenizer.pad_token_id = tokenizer.eos_token_id
+                  logger.warning("Tokenizer pad_token_id is None, using eos_token_id (%s) as pad_token_id.", tokenizer.eos_token_id)
              else:
-                 # Add a pad token if neither eos nor pad exists
+                 # Add a pad token if neither eos nor pad exists.
+                 # This *must* be done before loading the model or resizing embeddings.
+                 logger.warning("Tokenizer pad_token_id and eos_token_id are both None. Adding a [PAD] token.")
                  tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-                 model.resize_token_embeddings(len(tokenizer)) # Resize embeddings after adding token
                  tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids('[PAD]')
-                 logger.warning("Added and set [PAD] token, resized embeddings.")
+                 logger.info("Added new [PAD] token with ID %s.", tokenizer.pad_token_id)
+                 # Note: Resizing embeddings will be handled by the wrapper during initialization
+                 # if a compatible HF model instance is found.
+
+
+        # Load model
+        model = AutoModelForCausalLM.from_pretrained(model_id)
+
 
         # Instantiate the wrapper
-        # Simulate parameters that would come from the GUI
-        simulated_gen_config = GenerationConfig(
-            max_new_tokens=100,
-            temperature=0.8,
-            do_sample=True,
-            num_return_sequences=2, # Simulate asking for 2 chains
+        # Simulate parameters that would come from a GUI or config
+        # This GenerationConfig will override some defaults in the wrapper's base config for this call.
+        simulated_base_gen_config = GenerationConfig(
+            max_new_tokens=128,      # Limit generated tokens
+            temperature=0.85,        # Slightly higher temp for diversity in multiple chains
+            do_sample=True,          # Crucial for sampling-based generation
+            # num_return_sequences is intentionally NOT set here; it's set by the wrapper based on generate() argument
             pad_token_id=tokenizer.pad_token_id, # Pass pad_token_id explicitly
             eos_token_id=tokenizer.eos_token_id, # Pass eos_token_id explicitly
+            # Add other parameters based on tuning recommendations if desired
+            repetition_penalty=1.1 # Apply repetition penalty
         )
 
+        # Instantiate the wrapper, enabling self-consistency flags in init
+        # These flags inform the wrapper's default behavior if generate() args are None
         cot_wrapper = ChainOfThoughtWrapper(
             model=model,
             tokenizer=tokenizer,
-            generation_config=simulated_gen_config,
+            generation_config=simulated_base_gen_config, # Pass overrides here if desired as base
             device=device,
-            self_consistency=True, # Simulate SC enabled
-            consistency_rounds=2, # Simulate consistency rounds setting
+            self_consistency_enabled=True, # Simulate SC enabled
+            consistency_rounds=5, # Simulate consistency rounds setting
+            final_answer_tag="Final Answer:", # Use a slightly different tag for demo
+            # Keep factual emphasis on for demo
+            emphasize_factual=True,
+            allow_uncertainty_phrase="If you cannot determine a definitive answer, state that.",
         )
 
         # Prepare input prompt
-        prompt_text = "What is 2 + 2? Think step-by-step."
-        input_enc = tokenizer(prompt_text, return_tensors='pt').to(device)
-
+        # Use a prompt that encourages steps and a clear answer
+        prompt_text = "If a train travels at 60 mph for 2.5 hours, how far does it travel? Calculate step-by-step."
         logger.info(f"Generating reasoning for prompt: '{prompt_text}'")
 
         # Generate outputs
-        # The num_return_sequences from simulated_gen_config will be used here
+        # We explicitly pass num_return_sequences to the generate call (e.g., from GUI slider)
+        num_chains_to_generate = 3 # Simulate GUI setting num_chains slider to 3
+        logger.info(f"Calling wrapper.generate() requesting {num_chains_to_generate} chains.")
+
+        start_time = time.time()
         outputs = cot_wrapper.generate(
-            input_ids=input_enc['input_ids'],
-            attention_mask=input_enc['attention_mask']
+            input_text=prompt_text,
+            # No explicit generation_config override here; uses the base config initialized in the wrapper
+            # but you *could* pass overrides like: generation_config=GenerationConfig(temperature=1.0)
+            num_return_sequences=num_chains_to_generate, # Pass the desired number of sequences here
         )
+        end_time = time.time()
+        logger.info(f"Generation of {len(outputs.get('sequences', []))} sequences took {end_time - start_time:.2f} seconds.")
+
+
+        # --- Process Results (including simulated Self-Consistency voting logic) ---
+        print("\n" + "="*50)
+        print("--- Generation Results ---")
+        print("="*50)
+
+        full_texts = outputs.get('full_texts', [])
+        reasoning_steps = outputs.get('reasoning_steps', [])
+        final_answers_raw = outputs.get('final_answers', []) # Raw answers from wrapper
+
+        if not full_texts:
+             print("No chains were generated or parsed.")
+        else:
+            for i, (full_text, steps, final_answer_raw) in enumerate(zip(full_texts, reasoning_steps, final_answers_raw)):
+                print(f"\n--- Chain {i+1} ---")
+                print("Full Text (Cleaned):")
+                print(full_text)
+                print("\nReasoning Steps Parsed:")
+                if steps:
+                    # Ensure steps is a list before iterating
+                    steps = steps if isinstance(steps, list) else []
+                    for j, step in enumerate(steps):
+                         # Ensure step is a string before printing
+                         if isinstance(step, str) and step.strip():
+                            print(f"  Step {j+1}: {step.strip()}")
+                         elif not isinstance(step, str):
+                              print(f"  [Step {j+1} has invalid format]")
+                    if not steps: # If steps list was empty after checks
+                        print("  [No steps parsed]")
+                else: # If steps was None or not a list initially
+                    print("  [No steps parsed]")
+                print("\nFinal Answer Parsed (Raw):")
+                # Ensure raw answer is a string before printing
+                display_raw_answer = final_answer_raw if isinstance(final_answer_raw, str) and final_answer_raw.strip() else "[No final answer parsed]"
+                print(f"  '{display_raw_answer}'")
 
-        # Process results (including simulated Self-Consistency voting logic)
-        print("\n--- Generation Results ---")
-        for i, (full_text, steps, final_answer) in enumerate(zip(outputs['full_texts'], outputs['reasoning_steps'], outputs['final_answers'])):
-            print(f"\n--- Chain {i+1} ---")
-            print("Full Text:")
-            print(full_text)
-            print("\nReasoning Steps:")
-            if steps:
-                for j, step in enumerate(steps):
-                    print(f"  Step {j+1}: {step}")
-            else:
-                print("  [No steps parsed]")
-            print("\nFinal Answer:")
-            print(f"  {final_answer or '[No final answer parsed]'}")
 
         # --- Simulate Self-Consistency Voting (as would be done in GUI) ---
-        print("\n--- Self-Consistency Voting ---")
-        final_answers = [ans for ans in outputs['final_answers'] if ans.strip()] # Filter empty answers
-        if final_answers:
-            answer_counts = Counter(final_answers)
-            most_common_answer, count = answer_counts.most_common(1)[0]
-            print(f"Raw Answers Submitted for Voting: {final_answers}")
-            print(f"Answer Counts: {dict(answer_counts)}")
-            print(f"Consensus Answer: '{most_common_answer}' (Voted by {count} chain(s))")
+        print("\n" + "="*50)
+        print("--- Simple Self-Consistency Voting Simulation ---")
+        print("="*50)
+
+        if final_answers_raw:
+            # Perform the actual voting using the helper functions
+            consensus_answer, answer_distribution_dict = perform_self_consistency_voting(final_answers_raw)
+            answer_distribution = Counter(answer_distribution_dict) # Convert to Counter for display
+
+            print(f"Raw Answers Submitted for Voting: {final_answers_raw}")
+            print(f"Normalized Answers for Voting: {list(answer_distribution_dict.keys())}") # Show unique normalized answers
+            print(f"Answer Counts: {dict(answer_distribution)}")
+
+            if consensus_answer:
+                print(f"\nConsensus Answer: '{consensus_answer}'")
+                # Get count of the winning normalized answer
+                winner_count = answer_distribution.get(normalize_answer(consensus_answer), 0)
+                print(f"(Voted by {winner_count} chain(s) out of {len(final_answers_raw)})")
+
+                # Optional: Check for ties (more sophisticated tie-breaking would go here in a real voter)
+                if len(answer_distribution) > 1 and answer_distribution.most_common(2)[0][1] == answer_distribution.most_common(2)[1][1]:
+                     print("Note: There is a tie for the most common normalized answer.")
+
+            else:
+                print("No valid final answers found for voting.")
         else:
-            print("No valid final answers found for voting.")
+            print("No final answers were parsed from any chain for voting.")
 
 
     except Exception as e:
-        logger.error("Example usage failed: %s", e)
+        logger.error("An error occurred during the example usage: %s", e)
         import traceback
         traceback.print_exc() # Print detailed traceback for the example failure
 
-    print("\n--- Example Usage End ---")
+    print("\n--- Example Usage End ---")
+    # Attempt final cleanup
+    if torch.cuda.is_available():
+         try:
+             torch.cuda.empty_cache()
+             print("GPU memory cache cleared.")
+         except Exception as cleanup_e:
+              print(f"Error during final cuda empty_cache: {cleanup_e}")
+    gc.collect()
+    print("Garbage collected.")