import json import random from pathlib import Path import tiktoken from datasets import Dataset from opencompass.datasets.base import BaseDataset from opencompass.openicl import BaseEvaluator from opencompass.registry import LOAD_DATASET def get_unique_entries(file_path, n, language, unique_arg1=False, unique_arg2=False, unique_combination=False): seen_arg1 = set() seen_arg2 = set() seen_combinations = set() results = [] with open(file_path, 'r', encoding='utf-8') as file: lines = file.readlines() random.shuffle(lines) for line in lines: try: entry = json.loads(line.strip()) except json.JSONDecodeError: continue if entry.get('language') != language: continue key1 = entry.get('arg1', '') if unique_arg1 else '' key2 = entry.get('arg2', '') if unique_arg2 else '' combination = (key1, key2) if unique_combination else '' if (key1 not in seen_arg1 or not unique_arg1) and \ (key2 not in seen_arg2 or not unique_arg2) and \ (combination not in seen_combinations or not unique_combination): seen_arg1.add(key1) seen_arg2.add(key2) seen_combinations.add(combination) results.append(entry) if len(results) == n: break return results @LOAD_DATASET.register_module() class NeedleBenchParallelDataset(BaseDataset): @staticmethod def load( path: str, needle_file_name: str, length: int, depths: list[int], tokenizer_model: str, file_list: list[str], num_repeats_per_file: int, length_buffer: int, guide: bool, language: str, ): data = {'prompt': [], 'answer': []} tokenizer = tiktoken.encoding_for_model(tokenizer_model) files = Path(path).glob('*.jsonl') for file in files: if file.name == needle_file_name: needle_file_path = file predefined_needles_bak = get_unique_entries(needle_file_path, len(depths), language, unique_arg1=True, unique_arg2=True, unique_combination=True) def _generate_context(tokens_context, depths, needles): insertion_points = [ int(len(tokens_context) * (depth / 100)) for depth in depths ] cumulative_inserted_length = 0 for i, needle in enumerate(needles): needle_tokens = _get_tokens_from_context(needle) current_insertion_point = min( insertion_points[i] + cumulative_inserted_length, len(tokens_context)) tokens_context = tokens_context[:current_insertion_point] + \ needle_tokens + tokens_context[current_insertion_point:] cumulative_inserted_length += len(needle_tokens) new_context = _decode_tokens(tokens_context) return new_context def _get_tokens_from_context(context): if isinstance(context, list): return [tokenizer.encode(item) for item in context] else: return tokenizer.encode(context) def _decode_tokens(tokens): return tokenizer.decode(tokens) def _modify_retrieval_question(retrieval_question): if language == 'Chinese': parts = retrieval_question.split('请按照') guide_retrieval_question = (parts[0] + '在回答之前，请思考文档中与此问题' '最相关的内容是什么。请按照' + parts[1]) return guide_retrieval_question elif language == 'English': parts = retrieval_question.split('Please answer in the format') guide_retrieval_question = ( parts[0] + 'Before answering, please consider' ' what in the document is most relevant to this question.' ' Please answer in the format' + parts[1]) return guide_retrieval_question else: raise ValueError(f"Language '{language}' is not supported.") def _generate_prompt(context, retrieval_question): if guide: retrieval_question = _modify_retrieval_question( retrieval_question) if language == 'Chinese': prompt = ('你是一个善于回答用户问题的智能AI助手\n' '请保持你的回答简洁清楚。不要说和下面文档中的无关的话' '，或重复你的回答\n请先仔细阅读下面的文档再依次回答' f'最后提出的问题\n用户现在给你的文档是{context}\n\n' f'现在请问：{retrieval_question}\n') elif language == 'English': prompt = ( 'You are an intelligent AI assistant skilled in ' 'answering user questions.\n' 'Please keep your answers concise and clear. Do not' ' talk about irrelevant topics or repeat your ' 'answers.\n' f'The document given to you by the user is {context}' f'\n\nNow, the questions are: {retrieval_question}\n') else: raise ValueError(f"Language '{language}' is not supported.") return prompt files = Path(path).glob('*.jsonl') for file in files: if file.name not in file_list: continue with open(file, 'r', encoding='utf-8') as f: lines_bak = [json.loads(line.strip()) for line in f] lines = lines_bak.copy() for counter in range(num_repeats_per_file): random.seed(counter) random.shuffle(lines) predefined_needles = predefined_needles_bak.copy() random.shuffle(predefined_needles) needles = [ '\n' + item['needle'] + '\n' for item in predefined_needles ] keywords = [item['arg2'] for item in predefined_needles] if language == 'Chinese': questions = '、'.join([ item['retrieval_question'].split('？')[0] + '？' for item in predefined_needles ]) answers_format = '、'.join([ item['retrieval_question'].split("'")[1].split('。')[0] for item in predefined_needles ]) retrieval_question = questions + "请按照'" + \ answers_format + "'的格式回答。" elif language == 'English': questions = '、'.join([ item['retrieval_question'].split('?')[0] + '?' for item in predefined_needles ]) answers_format = '、'.join([ item['retrieval_question'].split("'")[1].split('.')[0] for item in predefined_needles ]) retrieval_question = questions + \ "Please answer in the format of '" + \ answers_format + "'" context_length = length - length_buffer target_length_per_record = context_length - \ sum(len(tokens) for tokens in _get_tokens_from_context(needles)) target_length_per_record = max(target_length_per_record, 0) accumulated_tokens = [] for line in lines: tokens_current_line = _get_tokens_from_context( line['text']) accumulated_tokens.extend(tokens_current_line) if len(accumulated_tokens) >= target_length_per_record: break processed_text = _generate_context( accumulated_tokens[:target_length_per_record], depths, needles) processed_prompt = _generate_prompt(processed_text, retrieval_question) data['prompt'].append(processed_prompt) data['answer'].append('*'.join(keywords) + '#' + '*'.join(map(str, depths))) dataset = Dataset.from_dict({ 'prompt': data['prompt'], 'answer': data['answer'], }) return dataset class NeedleBenchParallelEvaluator(BaseEvaluator): def levenshtein_distance(self, s1, s2): if len(s1) < len(s2): return self.levenshtein_distance(s2, s1) if len(s2) == 0: return len(s1) previous_row = range(len(s2) + 1) for i, c1 in enumerate(s1): current_row = [i + 1] for j, c2 in enumerate(s2): insertions = previous_row[j + 1] + 1 deletions = current_row[j] + 1 substitutions = previous_row[j] + (c1 != c2) current_row.append(min(insertions, deletions, substitutions)) previous_row = current_row return previous_row[-1] def score(self, predictions, gold): if len(predictions) != len(gold): return {'error': 'predictions and gold have different lengths'} print('predictions:', predictions) print('gold:', gold) details = [] depths = [int(i) for i in gold[0].split('#')[1].split('*')] scores_by_depth = {depth: 0 for depth in depths} for prediction, reference in zip(predictions, gold): print(reference) keywords = reference.split('#')[0].split('*') print(keywords) for keyword, depth in zip(keywords, depths): print('iterating:', keyword, depth) if keyword in prediction: print(f'{keyword} at depth {depth} is in {prediction}') scores_by_depth[depth] += 100 / (len(predictions)) average_score = sum(scores_by_depth.values()) / len(scores_by_depth) flattened_scores = { 'Depth' + str(depth): score for depth, score in scores_by_depth.items() } result = { **flattened_scores, 'details': details, 'average_score': average_score } return result