import torch from numba import njit from modules import shared def process_llamacpp_cache(model, new_sequence, past_sequence): if len(past_sequence) == 0 or len(new_sequence) == 0: return past_sequence i1, i2, j1, j2 = find_longest_common_substring_indices(past_sequence, new_sequence) overlap_length = i2 - i1 + 1 # Do StreamingLLM if i1 > 0 (ie the longest common subsequence is not a prefix) # and the overlap length is sufficiently long. if i1 > 0 and overlap_length > 0.2 * len(new_sequence): new_sequence = torch.tensor(new_sequence) past_sequence = torch.tensor(past_sequence) prefix_length = find_prefix_length(past_sequence[:i1], new_sequence[:j1]) sink_length = prefix_length if sink_length < shared.args.attention_sink_size: sink_length = shared.args.attention_sink_size removed_length = i1 - sink_length matching_prefix = past_sequence[:prefix_length] removed_chunk = past_sequence[sink_length:i1] overlapping_sequence = new_sequence[j1:j2 + 1] added_chunk = new_sequence[j2 + 1:] # print(past_sequence) # print(new_sequence) print() print('MATCHING PREFIX=', repr(shared.tokenizer.decode(matching_prefix))) print('ADDED CHUNK=', repr(shared.tokenizer.decode(added_chunk))) print('REMOVED CHUNK=', repr(shared.tokenizer.decode(removed_chunk))) print() # Remove interval [sink_length, sink_length + removed_length) from the context # Subtract removed_length from model.n_tokens model._ctx.kv_cache_seq_rm(0, sink_length, sink_length + removed_length) model._ctx.kv_cache_seq_shift(0, sink_length + removed_length, -1, -removed_length) new_sequence = new_sequence.tolist() model.input_ids[:j2 + 1] = new_sequence[:j2 + 1] model.n_tokens = j2 + 1 return new_sequence[:j2 + 1] else: return past_sequence def find_prefix_length(past_seq, seq_tensor): ''' Given two torch tensors, finds the length of the longest common prefix between the two. ''' min_length = min(past_seq.shape[0], seq_tensor.shape[0]) indices = torch.nonzero(~torch.eq(past_seq[:min_length], seq_tensor[:min_length])) if len(indices) > 0: prefix_length = indices[0].item() else: prefix_length = min_length return prefix_length @njit def find_longest_common_substring_indices(list1, list2): ''' Given two lists, solves the Longest Common Substring problem. It returns the indices where the substring starts and ends in s1 and s2. Example: ir, jr, ir2, jr2 = find_longest_common_substring_indices(s1, s2) print(s1[ir:jr + 1]) print(s2[ir2:jr2 + 1]) Adapted from https://rosettacode.org/wiki/Longest_common_substring#Python ''' len_list1, len_list2 = len(list1), len(list2) start_index_list1, end_index_list1 = 0, -1 start_index_list2, end_index_list2 = 0, -1 # for index1 in tqdm(range(0, len_list1), desc="StreamingLLM prompt comparison", leave=False): for index1 in range(0, len_list1): try: index2 = list2.index(list1[index1]) except: continue while index2 >= 0: temp_index1, temp_index2 = index1, index2 while temp_index1 < len_list1 and temp_index2 < len_list2 and list2[temp_index2] == list1[temp_index1]: if temp_index1 - index1 >= end_index_list1 - start_index_list1: start_index_list1, end_index_list1 = index1, temp_index1 start_index_list2, end_index_list2 = index2, temp_index2 temp_index1 += 1 temp_index2 += 1 try: index2 = list2.index(list1[index1], index2 + 1) except: break return start_index_list1, end_index_list1, start_index_list2, end_index_list2