back to last promising version with treebuild crashlog, revising whole tree gen thing with Claude 5

app.py

@@ -19,6 +19,14 @@ from graphviz import Digraph
 from PIL import Image, ImageDraw, ImageFont
 from io import BytesIO
 import functools
+import logging
+
+# Set up custom logger
+custom_logger = logging.getLogger("custom_logger")
+custom_logger.setLevel(logging.INFO)
+custom_handler = logging.StreamHandler()
+custom_handler.setFormatter(logging.Formatter('%(message)s'))
+custom_logger.addHandler(custom_handler)
 
 # Load environment variables
 load_dotenv()
@@ -203,13 +211,16 @@ def produce_next_token_ids(input_ids, model, topk, sub_token_id):
 
 
 def build_def_tree(input_ids, data, base_prompt, model, tokenizer, config, depth=0, max_depth=25, cumulative_prob=1.0):
-    logger.info(f"build_def_tree called with depth={depth}, cumulative_prob={cumulative_prob:.4f}")
     if depth >= max_depth or cumulative_prob < config.CUTOFF:
         return
 
     current_prompt = tokenizer.decode(input_ids[0], skip_special_tokens=True)
     
-    yield f"\nDepth {depth}: {current_prompt}      PROB: {cumulative_prob:.4f}"
+    # Extract only the part that extends the base prompt
+    extended_prompt = current_prompt[len(base_prompt):].strip()
+    
+    # Use a custom logger to avoid the "INFO:__main__:" prefix
+    custom_logger.info(f"Depth {depth}: {extended_prompt}      PROB: {cumulative_prob:.4f}")
     
     top_k_ids, top_k_probs = produce_next_token_ids(input_ids, model, config.TOPK, config.SUB_TOKEN_ID)
 
@@ -238,6 +249,7 @@ def build_def_tree(input_ids, data, base_prompt, model, tokenizer, config, depth
         yield from build_def_tree(new_input_ids, new_child, base_prompt, model, tokenizer, config, depth=depth+1, max_depth=max_depth, cumulative_prob=new_cumulative_prob)
 
 
+
 def generate_definition_tree(base_prompt, embedding, model, tokenizer, config):
     logger.info(f"Starting generate_definition_tree with base_prompt: {base_prompt}")
     results_dict = {"token": "", "cumulative_prob": 1, "children": []}
@@ -421,10 +433,8 @@ def process_input(selected_sae, feature_number, weight_type, use_token_centroid,
             for item in tree_generator:
                 if isinstance(item, str):
                     log_output.append(item)
-                    logger.info(f"Tree generation step: {item}")  # Log each step
                 else:
                     tree_data = item
-                    logger.info("Received tree data")
             
             # Join the log output into a single string
             log_text = "\n".join(log_output)
@@ -440,6 +450,7 @@ def process_input(selected_sae, feature_number, weight_type, use_token_centroid,
                 logger.error("Failed to generate tree data")
                 return "Error: Failed to generate tree data.", None
         
+        
         return "Mode not recognized or not implemented in this step.", None
     
     except Exception as e: