Update README.md

README.md

@@ -245,14 +245,12 @@ def chunk_text_with_max_chunk_size(model, text, tokenizer, prob_threshold=0.5,ma
         
         unchunk_tokens = 0
         backup_pos = None
-        best_logits = torch.finfo(torch.float32).min 
-        is_chunk_start = True
-        
-        STEP = round(((MAX_TOKENS - 2)//2 )*1.75) #(MAX_TOKENS - 2)//2
+        best_logits = torch.finfo(torch.float32).min         
+        STEP = round(((MAX_TOKENS - 2)//2)*1.75 )
         print(f"Processing {input_ids.shape[1]} tokens...")
-        while windows_end <= input_ids.shape[1]:
-            
-            windows_end = windows_start + MAX_TOKENS - 2
+        # while windows_end <= input_ids.shape[1]:#记得改成windstart
+        while windows_start < input_ids.shape[1]:#记得改成windstart    
+            windows_end = windows_start + MAX_TOKENS - 2 
             ids = torch.cat((CLS, input_ids[:, windows_start:windows_end], SEP), 1)
             ids = ids.to(model.device)
             output = model(
@@ -279,10 +277,7 @@ def chunk_text_with_max_chunk_size(model, text, tokenizer, prob_threshold=0.5,ma
                 # manually chunk
                 if unchunk_tokens + unchunk_tokens_this_window > max_tokens_per_chunk:
                     big_windows_end = max_tokens_per_chunk - unchunk_tokens
-                    if is_chunk_start:
-                        max_value, max_index= logit_diff[:,1:big_windows_end].max(),  logit_diff[:,1:big_windows_end].argmax() + 1
-                    else:
-                        max_value, max_index= logit_diff[:,1:big_windows_end].max(),  logit_diff[:,1:big_windows_end].argmax() + 1
+                    max_value, max_index= logit_diff[:,1:big_windows_end].max(),  logit_diff[:,1:big_windows_end].argmax() + 1
                     if best_logits < max_value:
                         backup_pos = windows_start + max_index
                     
@@ -295,15 +290,17 @@ def chunk_text_with_max_chunk_size(model, text, tokenizer, prob_threshold=0.5,ma
                     best_logits = torch.finfo(torch.float32).min
                     backup_pos = -1
                     unchunk_tokens = 0
-                    is_chunk_start = True
                     
                 # auto chunk    
                 else:
+                    
                     if len(greater_rows_indices) >= 2:
                         for gi, (gri0,gri1) in enumerate(zip(greater_rows_indices[:-1],greater_rows_indices[1:])):
+                            
                             if gri1 - gri0 > max_tokens_per_chunk:
                                 greater_rows_indices=greater_rows_indices[:gi+1]
                                 break
+                                
                     split_str_pos = [tokens.token_to_chars(sp + windows_start + 1).start for sp in greater_rows_indices if sp > 0]
                     split_str_poses = split_str_poses + split_str_pos
                     token_pos = token_pos+ [sp + windows_start for sp in greater_rows_indices if sp > 0]
@@ -312,20 +309,20 @@ def chunk_text_with_max_chunk_size(model, text, tokenizer, prob_threshold=0.5,ma
                     best_logits = torch.finfo(torch.float32).min
                     backup_pos = -1
                     unchunk_tokens = 0
-                    is_chunk_start = True
 
             else:
 
-                unchunk_tokens_this_window = min(windows_end - windows_start,STEP)
+                # unchunk_tokens_this_window = min(windows_end - windows_start,STEP)
+                unchunk_tokens_this_window = min(windows_start+STEP,input_ids.shape[1]) - windows_start
+
                 # manually chunk
                 if unchunk_tokens + unchunk_tokens_this_window > max_tokens_per_chunk:
                     big_windows_end =  max_tokens_per_chunk - unchunk_tokens
-                    if is_chunk_start:
-                        max_value, max_index= logit_diff[:,1:big_windows_end].max(),  logit_diff[:,1:big_windows_end].argmax() + 1
-                    else:
+                    if logit_diff.shape[1] > 1:
+                        
                         max_value, max_index= logit_diff[:,1:big_windows_end].max(),  logit_diff[:,1:big_windows_end].argmax() + 1
-                    if best_logits < max_value:
-                        backup_pos = windows_start + max_index
+                        if best_logits < max_value:
+                            backup_pos = windows_start + max_index
                         
                         
                     windows_start = backup_pos
@@ -335,21 +332,16 @@ def chunk_text_with_max_chunk_size(model, text, tokenizer, prob_threshold=0.5,ma
                     best_logits = torch.finfo(torch.float32).min
                     backup_pos = -1
                     unchunk_tokens = 0
-                    is_chunk_start = True
                 else:
                 # auto leave
-                    if is_chunk_start:
+                    if logit_diff.shape[1] > 1:
                         max_value, max_index= logit_diff[:,1:].max(),  logit_diff[:,1:].argmax() + 1
-
-                    else:
-                            max_value, max_index= logit_diff[:,1:].max(),  logit_diff[:,1:].argmax() + 1
-                    if best_logits < max_value:
-                        best_logits = max_value
-                        backup_pos = windows_start + max_index
+                        if best_logits < max_value:
+                            best_logits = max_value
+                            backup_pos = windows_start + max_index
                        
                     unchunk_tokens = unchunk_tokens + STEP
                     windows_start = windows_start + STEP
-                    is_chunk_start = False
 
         substrings = [
             text[i:j] for i, j in zip([0] + split_str_poses, split_str_poses + [len(text)])