next attempt at tree functionality 7

app.py

@@ -195,22 +195,31 @@ def produce_next_token_ids(input_ids, model, topk, sub_token_id):
         outputs = model(input_ids)
         logits = outputs.logits
 
-    if logits.size(1) == 0:  # Check if there are logits to process
-        logger.error("Logits are empty. Cannot produce next token IDs.")
+    # Check if logits is empty or invalid
+    if logits is None or logits.size(1) == 0:
+        logger.error("Logits are empty or invalid.")
         return None, None
 
     last_logits = logits[:, -1, :]
 
-    if last_logits.size(0) == 0 or last_logits.size(1) == 0:  # Check if last logits are valid
-        logger.error("Last logits are empty. Cannot produce next token IDs.")
+    # Check if last_logits is empty or invalid
+    if last_logits.size(0) == 0 or last_logits.size(1) == 0:
+        logger.error("Last logits are empty or invalid.")
         return None, None
 
     last_logits[:, sub_token_id] = float('-inf')
     softmax_probs = torch.softmax(last_logits, dim=-1)
+    
+    # Check if softmax_probs are empty or invalid
+    if softmax_probs is None or softmax_probs.size(0) == 0:
+        logger.error("Softmax probabilities are empty or invalid.")
+        return None, None
+
     top_k_probs, top_k_ids = torch.topk(softmax_probs, k=topk, dim=-1)
 
-    if top_k_ids.size(0) == 0 or top_k_probs.size(0) == 0:  # Check if we successfully got top-k IDs and probabilities
-        logger.error("Top-k IDs or probabilities are empty. Cannot produce next token IDs.")
+    # Check if top_k_ids and top_k_probs are valid
+    if top_k_ids.size(0) == 0 or top_k_probs.size(0) == 0:
+        logger.error("Top-k IDs or probabilities are empty or invalid.")
         return None, None
 
     return top_k_ids[0], top_k_probs[0]
@@ -225,7 +234,7 @@ def build_def_tree(input_ids, data, base_prompt, model, tokenizer, config, depth
     
     top_k_ids, top_k_probs = produce_next_token_ids(input_ids, model, config.TOPK, config.SUB_TOKEN_ID)
     
-    if top_k_ids is None or top_k_probs is None:  # Ensure that top_k_ids and top_k_probs are valid before proceeding
+    if top_k_ids is None or top_k_probs is None:
         logger.error(f"Failed to generate next token IDs at depth {depth}.")
         return
 
@@ -233,22 +242,18 @@ def build_def_tree(input_ids, data, base_prompt, model, tokenizer, config, depth
         if token_id == config.SUB_TOKEN_ID:
             continue  # Skip the substitute token to avoid circular definitions
 
-        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
+        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)
 
         new_cumulative_prob = cumulative_prob * top_k_probs[idx].item()
 
         if new_cumulative_prob < config.CUTOFF:
             continue
 
-        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)}")
+        token_str = tokenizer.decode([token_id], skip_special_tokens=True)
+        
+        if token_str is None or token_str == "":
+            logger.error(f"Token string is empty or invalid at depth {depth} for token ID {token_id}.")
             continue
 
         new_child = {
@@ -279,8 +284,6 @@ def generate_definition_tree(base_prompt, embedding, model, tokenizer, config, p
 
 
 
-
-
 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: