| import json |
| import os |
| import random |
| import re |
| 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, TEXT_POSTPROCESSORS |
|
|
|
|
| def get_random_line_by_language(file_path, language): |
| with open(file_path, 'r', encoding='utf-8') as file: |
| lines = [ |
| json.loads(line.strip()) for line in file |
| if json.loads(line.strip())['language'] == language |
| ] |
|
|
| if lines: |
| random_line = random.choice(lines) |
| return { |
| 'needle': random_line['needle'], |
| 'retrieval_question': random_line['retrieval_question'], |
| 'keyword': random_line['arg2'] |
| } |
| else: |
| return None |
|
|
|
|
| @LOAD_DATASET.register_module() |
| class NeedleBenchOriginDataset(BaseDataset): |
|
|
| @staticmethod |
| def load( |
| path: str, |
| length: int, |
| depth: int, |
| tokenizer_model: str, |
| file_list: list[str], |
| num_repeats_per_file: int, |
| length_buffer: int, |
| guide: bool, |
| language: str, |
| needle_file_name: str, |
| ): |
| data = {'prompt': [], 'answer': []} |
| tokenizer = tiktoken.encoding_for_model(tokenizer_model) |
|
|
| def _generate_context(tokens_context, depth_percent, needle): |
| tokens_needle = _get_tokens_from_context(needle) |
| insertion_point = int(len(tokens_context) * (depth_percent / 100)) |
| tokens_context = (tokens_context[:insertion_point] + |
| tokens_needle + tokens_context[insertion_point:]) |
| new_context = _decode_tokens(tokens_context) |
| return new_context |
|
|
| def _get_tokens_from_context(context): |
| 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'用户现在给你的文档是{context}\n\n' |
| f'现在请问：{retrieval_question}') |
| 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 question is: {retrieval_question}') |
| 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) |
| needle_file_path = os.path.join(path, needle_file_name) |
| random_needle = get_random_line_by_language( |
| needle_file_path, language) |
| needle = '\n' + random_needle['needle'] + '\n' |
| retrieval_question = random_needle['retrieval_question'] |
| keyword = random_needle['keyword'] |
|
|
| context_length = length - length_buffer |
| target_length_per_record = context_length - len( |
| _get_tokens_from_context(needle)) |
| 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], depth, |
| needle) |
|
|
| processed_prompt = _generate_prompt(processed_text, |
| retrieval_question) |
|
|
| data['prompt'].append(processed_prompt) |
| data['answer'].append(needle + '*' + keyword) |
|
|
| dataset = Dataset.from_dict({ |
| 'prompt': data['prompt'], |
| 'answer': data['answer'], |
| }) |
| return dataset |
|
|
|
|
| class NeedleBenchOriginEvaluator(BaseEvaluator): |
|
|
| def __init__(self, use_trim=False): |
| self.use_trim = use_trim |
|
|
| @staticmethod |
| def _trim_prediction(prediction, reference): |
| """Trims the prediction string based on the length of the reference |
| string. |
| |
| Args: |
| prediction (str): The prediction string. |
| reference (str): The reference string. |
| |
| Returns: |
| str: The trimmed prediction string. |
| """ |
| l08 = int(0.8 * len(reference)) |
| l12 = int(1.2 * len(reference)) |
| trimmed_prediction = prediction[:l12] |
|
|
| if len(trimmed_prediction) > l08 and \ |
| reference[-1] in trimmed_prediction[l08:]: |
| end_pos = l08 + trimmed_prediction[l08:].index(reference[-1]) + 1 |
| trimmed_prediction = trimmed_prediction[:end_pos] |
|
|
| return trimmed_prediction |
|
|
| 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'} |
|
|
| total_score = 0 |
| details = [] |
| for prediction, reference in zip(predictions, gold): |
| keyword = reference.split('*')[1] |
| reference = reference.split('*')[0] |
| raw_prediction = prediction |
| prediction = re.sub(r'\s+', '', prediction) |
| reference = re.sub(r'\s+', '', reference) |
|
|
| if self.use_trim: |
| prediction = NeedleBenchOriginEvaluator._trim_prediction( |
| prediction, reference) |
|
|
| edit_distance = self.levenshtein_distance(prediction, reference) |
| max_len = max(len(prediction), len(reference)) |
| score = 100 * (1 - |
| edit_distance / max_len) if max_len != 0 else 100 |
|
|
| if keyword in raw_prediction: |
| print(f'{keyword} is in {prediction}') |
| score = 100 |
| else: |
| print(f'{keyword} is not in {prediction}') |
| score = 0.2 * score |
|
|
| detail = { |
| 'pred': prediction, |
| 'answer': reference, |
| 'edit_distance': edit_distance, |
| 'score': score |
| } |
| total_score += score |
| details.append(detail) |
|
|
| average_score = total_score / len(predictions) if predictions else 0 |
| result = {'score': average_score, 'details': details} |
| return result |
|
|
|
|
| @TEXT_POSTPROCESSORS.register_module('needlebench') |
| def needlebench_postprocess(text: str) -> str: |
| return text |
|
|
|
|
| @TEXT_POSTPROCESSORS.register_module('needlebench_dataset') |
| def needlebench_dataset_postprocess(text: str) -> str: |
| return text |
|
|
