next attempt at tree functionality 5

app.py

@@ -208,21 +208,34 @@ def build_def_tree(input_ids, data, base_prompt, model, tokenizer, config, depth
     current_prompt = tokenizer.decode(input_ids[0], skip_special_tokens=True)
     if progress_callback:
         progress_callback(f"Depth {depth}: {current_prompt}      PROB: {cumulative_prob}\n")
-    top_k_ids, top_k_probs = produce_next_token_ids(input_ids, model, config.TOPK, config.SUB_TOKEN_ID)
+    
+    try:
+        top_k_ids, top_k_probs = produce_next_token_ids(input_ids, model, config.TOPK, config.SUB_TOKEN_ID)
+    except Exception as e:
+        logger.error(f"Error generating next token IDs at depth {depth}: {str(e)}")
+        return
 
     for idx, token_id in enumerate(top_k_ids.tolist()):
         if token_id == config.SUB_TOKEN_ID:
-            continue
+            continue  # Skip the substitute token to avoid circular definitions
 
-        token_id_tensor = torch.tensor([token_id], dtype=torch.long).to(model.device)
-        new_input_ids = torch.cat([input_ids, token_id_tensor.view(1, 1)], dim=-1)
+        try:
+            token_id_tensor = torch.tensor([token_id], dtype=torch.long).to(model.device)
+            new_input_ids = torch.cat([input_ids, token_id_tensor.view(1, 1)], dim=-1)
+        except IndexError as e:
+            logger.error(f"IndexError in processing token ID {token_id} at depth {depth}: {str(e)}")
+            continue
 
         new_cumulative_prob = cumulative_prob * top_k_probs[idx].item()
 
         if new_cumulative_prob < config.CUTOFF:
             continue
 
-        token_str = tokenizer.decode([token_id], skip_special_tokens=True)
+        try:
+            token_str = tokenizer.decode([token_id], skip_special_tokens=True)
+        except Exception as e:
+            logger.error(f"Error decoding token ID {token_id} at depth {depth}: {str(e)}")
+            continue
 
         new_child = {
             "token_id": token_id,
@@ -235,14 +248,14 @@ def build_def_tree(input_ids, data, base_prompt, model, tokenizer, config, depth
         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, progress_callback=progress_callback)
 
 
-
-
 def generate_definition_tree(base_prompt, embedding, model, tokenizer, config, progress_callback=None):
     results_dict = {"token": "", "cumulative_prob": 1, "children": []}
 
+    # Reset the token embedding
     token_embedding = torch.unsqueeze(embedding, dim=0).to(model.device)
     update_token_embedding(model, config.SUB_TOKEN_ID, token_embedding)
 
+    # Clear the model's cache if it has one
     if hasattr(model, 'reset_cache'):
         model.reset_cache()
 
@@ -251,6 +264,9 @@ def generate_definition_tree(base_prompt, embedding, model, tokenizer, config, p
 
     return results_dict
 
+
+
+
 def find_max_min_cumulative_weight(node, current_max=0, current_min=float('inf')):
     current_max = max(current_max, node.get('cumulative_prob', 0))
     if node.get('cumulative_prob', 1) > 0: