import argparse import json import os import shutil from pathlib import Path import matplotlib.pyplot as plt import numpy as np import pandas as pd import seaborn as sns import tiktoken from matplotlib.colors import LinearSegmentedColormap class CDMEDatasetProcessor: def __init__(self, path, output_path, tokenizer_model='gpt-4', num_records_per_file=10, length_buffer=200, guided=False, file_list=[]): self.path = path self.output_path = output_path self.tokenizer = tiktoken.encoding_for_model(tokenizer_model) self.num_records_per_file = num_records_per_file self.length_buffer = length_buffer self.guided = guided self.file_list = file_list def process_files(self, context_lengths, needle, retrieval_question, document_depth_percent_intervals, document_depth_percent_interval_type='linear'): files = Path(self.path).glob('*.jsonl') for file in files: if os.path.basename(file) in self.file_list: self.process_file(file, context_lengths, needle, retrieval_question, document_depth_percent_intervals, document_depth_percent_interval_type) def process_file(self, file, context_lengths, needle, retrieval_question, document_depth_percent_intervals, document_depth_percent_interval_type): with open(file, 'r', encoding='utf-8') as f: lines = [json.loads(line.strip()) for line in f] # set output file output_file = Path(self.output_path) / f'{file.stem}_processed.jsonl' output_file.parent.mkdir(parents=True, exist_ok=True) with open(output_file, 'w', encoding='utf-8') as out_f: for original_context_length in context_lengths: context_length = original_context_length - self.length_buffer target_length_per_record = context_length - len( self._get_tokens_from_context(needle)) for depth_percent in self._generate_depth_percents( document_depth_percent_intervals, document_depth_percent_interval_type): counter = 0 accumulated_tokens = [] for line in lines: tokens_current_line = self._get_tokens_from_context( line['text']) accumulated_tokens.extend(tokens_current_line) if len(accumulated_tokens) >= target_length_per_record: processed_text = self._generate_context( accumulated_tokens[:target_length_per_record], depth_percent, needle) processed_prompt = self._generate_prompt( processed_text, retrieval_question) json.dump( { 'prompt': processed_prompt, 'answer': needle, 'length': original_context_length, 'depth': int(depth_percent), }, out_f, ensure_ascii=False) out_f.write('\n') counter += 1 if counter >= self.num_records_per_file: break # reset accumulated_tokens for next record accumulated_tokens = [] def _generate_context(self, tokens_context, depth_percent, needle): tokens_needle = self._get_tokens_from_context(needle) # Insert the needle into the context at the specified depth percent insertion_point = int(len(tokens_context) * (depth_percent / 100)) tokens_context = (tokens_context[:insertion_point] + tokens_needle + tokens_context[insertion_point:]) # Decode the tokens back to text new_context = self._decode_tokens(tokens_context) return new_context def _get_tokens_from_context(self, context): return self.tokenizer.encode(context) def _decode_tokens(self, tokens): return self.tokenizer.decode(tokens) def _generate_prompt(self, context, retrieval_question): if self.guided: prompt = ('你是一个善于回答用户问题的智能AI助手\n' '请保持你的回答简洁清楚。不要说和下面文档中的无关的话，或重复你的回答\n' f'用户现在给你的文档是{context}\n\n' f'现在请问：{retrieval_question}' f'提示：文档中与该问题最相关的句子是_______') else: prompt = ('你是一个善于回答用户问题的智能AI助手\n' '请保持你的回答简洁清楚。不要说和下面文档中的无关的话，或重复你的回答\n' f'用户现在给你的文档是{context}\n\n' f'现在请问：{retrieval_question}') return prompt def _generate_depth_percents(self, intervals, interval_type): if interval_type == 'linear': return np.linspace(0, 100, num=intervals) elif interval_type == 'sigmoid': return [self._logistic(x) for x in np.linspace(0, 100, intervals)] else: raise ValueError('Unsupported interval type') @staticmethod def _logistic(x, L=100, x0=50, k=0.1): return np.round(L / (1 + np.exp(-k * (x - x0))), 3) class CDMEDataset(): @staticmethod def generate(processed_datasets_path, data_path, tokenizer_model, num_records_per_file, length_buffer, guided, file_list, context_lengths, needle, retrieval_question, document_depth_percent_intervals): # Check if the processed datasets directory exists if os.path.exists(processed_datasets_path): shutil.rmtree(processed_datasets_path) print('The existing processed datasets directory ' f'{processed_datasets_path} has been ' 'removed for a fresh start.') else: print('No existing processed datasets directory found at' f' {processed_datasets_path}. ' 'Starting with a fresh directory.') processor = CDMEDatasetProcessor( path=data_path, output_path=processed_datasets_path, tokenizer_model=tokenizer_model, num_records_per_file=num_records_per_file, length_buffer=length_buffer, guided=guided, file_list=file_list) processor.process_files(context_lengths, needle, retrieval_question, document_depth_percent_intervals) print('Datasets has been created.') @staticmethod def visualize(csv_file_paths): for file_path in csv_file_paths: df = pd.read_csv(file_path) # Split 'dataset' column to # get 'Context Length' and 'Document Depth' df['Context Length'] = df['dataset'].apply( lambda x: int(x.split('Length')[1].split('Depth')[0])) df['Document Depth'] = df['dataset'].apply( lambda x: float(x.split('Depth')[1].split('_')[0])) # Exclude 'Context Length' and 'Document Depth' columns model_columns = [ col for col in df.columns if col not in ['Context Length', 'Document Depth'] ] for model_name in model_columns[4:]: model_df = df[['Document Depth', 'Context Length', model_name]].copy() model_df.rename(columns={model_name: 'Score'}, inplace=True) # Create pivot table pivot_table = pd.pivot_table(model_df, values='Score', index=['Document Depth'], columns=['Context Length'], aggfunc='mean') # Calculate mean scores mean_scores = pivot_table.mean().values # Calculate overall score overall_score = mean_scores.mean() # Create heatmap and line plot plt.figure(figsize=(17.5, 8)) ax = plt.gca() cmap = LinearSegmentedColormap.from_list( 'custom_cmap', ['#F0496E', '#EBB839', '#0CD79F']) # Draw heatmap sns.heatmap(pivot_table, cmap=cmap, ax=ax, cbar_kws={'label': 'Score'}, vmin=0, vmax=100) # Set line plot data x_data = [i + 0.5 for i in range(len(mean_scores))] y_data = mean_scores # Create twin axis for line plot ax2 = ax.twinx() # Draw line plot ax2.plot(x_data, y_data, color='white', marker='o', linestyle='-', linewidth=2, markersize=8, label='Average Depth Score') # Set y-axis range ax2.set_ylim(0, 100) # Hide original y-axis ticks and labels ax2.set_yticklabels([]) ax2.set_yticks([]) # Add legend ax2.legend(loc='upper left') # Set chart title and labels ax.set_title(f'{model_name} 8K Context Performance\n' + 'Fact Retrieval Across Context Lengths ' + '("Needle In A Haystack")') ax.set_xlabel('Token Limit') ax.set_ylabel('Depth Percent') ax.set_xticklabels(pivot_table.columns.values, rotation=45) ax.set_yticklabels(pivot_table.index.values, rotation=0) # Add overall score as a subtitle plt.text(0.5, -0.13, f'Overall Score for {model_name}: ' f'{overall_score:.2f}', ha='center', va='center', transform=ax.transAxes, fontsize=13) # Save heatmap as PNG png_file_path = file_path.replace('.csv', f'_{model_name}.png') # plt.tight_layout() plt.savefig(png_file_path) plt.show() plt.close() # Close figure to prevent memory leaks # Print saved PNG file path print(f'Heatmap for {model_name} saved as: {png_file_path}') def main(): parser = argparse.ArgumentParser(description='Generate CDMEDataset.') parser.add_argument('--processed_datasets_path', type=str, default='./data/CDME/processed') parser.add_argument('--data_path', type=str, default='./data/CDME') parser.add_argument('--tokenizer_model', type=str, default='gpt-4') parser.add_argument('--num_records_per_file', type=int, default=10) parser.add_argument('--length_buffer', type=int, default=200) parser.add_argument('--guided', type=bool, default=False) parser.add_argument('--file_list', nargs='*', default=['zh_finance.jsonl']) parser.add_argument('--context_lengths', nargs='*', type=int, default=list(range(1000, 9000, 1000))) parser.add_argument('--needle', type=str, default='\n小明最喜欢的实习的地点就是上海人工智能实验室。\n') parser.add_argument('--retrieval_question', type=str, default='小明最喜欢的实习地点是哪里？' '你的回答格式应该为“小明最喜欢的实习地点就是________。”') parser.add_argument('--document_depth_percent_intervals', type=int, default=35) parser.add_argument('--plot', action='store_true', help='Visualize the dataset results') parser.add_argument('--csv_file_paths', nargs='*', default=['path/to/your/result.csv'], help='Paths to CSV files for visualization') args = parser.parse_args() if args.plot: if not args.csv_file_paths: print("Error: '--csv_file_paths' is required for visualization.") exit(1) CDMEDataset.visualize(args.csv_file_paths) else: doc_depth_intervals = args.document_depth_percent_intervals CDMEDataset.generate( processed_datasets_path=args.processed_datasets_path, data_path=args.data_path, tokenizer_model=args.tokenizer_model, num_records_per_file=args.num_records_per_file, length_buffer=args.length_buffer, guided=args.guided, file_list=args.file_list, context_lengths=args.context_lengths, needle=args.needle, retrieval_question=args.retrieval_question, document_depth_percent_intervals=doc_depth_intervals) if __name__ == '__main__': main()