InternBootcamp/examples/pipelines/cipher_data_generator.py

import argparse
import time
import shutil
import pandas as pd

from tqdm import tqdm

from internbootcamp.libs.cipher import *
from internbootcamp.bootcamp_utils import sample_sentences_by_percentage_distribution as sspd
from internbootcamp.bootcamp_utils.cipher_data_translator import translate
from internbootcamp.bootcamp_utils.cipher_prompt_enhance import enhance
from internbootcamp.bootcamp_utils.formatted_time import formatted_time

import hashlib
from internbootcamp.bootcamp_utils.random_things import random_word
from internbootcamp.bootcamp_utils import append_json_lines,get_difference_optimized,deduplicate_jsonl_by_field
import json
import random
import os

from tqdm import tqdm
class CipherDataGenerator():
    def __init__(self, env_class:BaseCipherEnvironment, input_list):
        """
        Initialize the Star Battle data generator with a specific grid size and attempts intervals.

        Args:
            size (int): The size of the Star Battle grid.
            attempts_intervals (list of tuples): A list of tuples defining the intervals for attempts.
                Each tuple contains two elements: the lower and upper bound of an interval.
                e.g., [(0, 5), (6, 10)] would create files for puzzles solved in 0-5 and 6-10 attempts.
        """
        super().__init__()
        self.env = env_class()
        self.input_list = input_list
        
    def generate_sample(self,plaintext: str = "", **kwargs):
        hashable_kwargs = json.dumps(kwargs, sort_keys=True)
        combined_str = plaintext + hashable_kwargs
        sample_id = hashlib.sha256(combined_str.encode('utf-8')).hexdigest()
        gen_steps,ciphertext = self.env.generator(plaintext, **kwargs)
        solve_steps,decodetext = self.env.solver(**kwargs)
        self.env.reset()
        
        return {
            "id": sample_id,
            "cipher_source": self.env.cipher_name,
            "plain": plaintext,
            "extra_args": kwargs,
            "encode_text": ciphertext,
            "decode_text": decodetext,
            "encode_rule": self.env.get_encode_rule(),
            "decode_rule": self.env.get_decode_rule(),
            "encode_steps": gen_steps,
            "decode_steps": solve_steps
        }
    
    def write_to_jsonl(self, timetamp):
        
        origin_data = os.path.join(f"examples/bootcamp_generator_outputs/{timetamp}/cipher//origin_data/", f"{self.env.cipher_name}.jsonl")
        origin_tmp_data = os.path.join(f"examples/bootcamp_generator_outputs/{timetamp}/cipher/origin_tmp_data/", f"{self.env.cipher_name}.jsonl")
        
        if not os.path.exists(f"examples/bootcamp_generator_outputs/{timetamp}/cipher/origin_data/"):
            # print(f"WARNING: examples/bootcamp_generator_outputs/{timetamp}/cipher/origin_data/ does not exist. Creating it.")
            os.makedirs(f"examples/bootcamp_generator_outputs/{timetamp}/cipher/origin_data/")
        
        if not os.path.exists(f"examples/bootcamp_generator_outputs/{timetamp}/cipher/origin_tmp_data/"):
            os.makedirs(f"examples/bootcamp_generator_outputs/{timetamp}/cipher/origin_tmp_data/")
        
        with open(origin_tmp_data, "w") as file:
            for sentence in self.input_list:
                puzzles_extra = {
                    'FourSquareCipher': {'str1': random_word(), 'str2': random_word()},
                    'Caesar_Cipher': {'shift': random.randint(1,26)},
                    'Kor_rule8_PortaCipher': {'key': random_word()},
                    'Autokey_Cipher': {'key': random_word()},
                    'Beale_Ciphers': {'codebook': "A big cat danced elegantly, finding giraffes happily in joyful kingdoms. Lions moved near owls, peacefully quizzing rabbits. Suddenly, tigers under vast waterfalls xeroxed yellow zebras."},
                    'Beaufort_Cipher': {'keyword': random_word()},
                    'KEYWORD_Cipher': {'keyword': random_word()},
                    'Rail_Fence_Cipher': {'num_rails': random.randint(2, 5)}                   
                }
                try:
                    if self.env.cipher_name in puzzles_extra:
                        extra_args = puzzles_extra[self.env.cipher_name]
                    else:
                        extra_args = {}
                    sample = self.generate_sample(plaintext=sentence.strip(),**extra_args)
                    if sample:
                        file.write(json.dumps(sample,ensure_ascii=False) + '\n')
                except Exception as e:
                    # print(f'{self.env.cipher_name} 错误：', e)
                    pass
        get_difference_optimized(file1=origin_tmp_data, file2=origin_data, key_field='id', output_file=origin_tmp_data)
        append_json_lines(source_file=origin_tmp_data, target_file=origin_data)

import json

def restructure_jsonl_file(file_path):
    """
    对指定的 JSONL 文件中的每一行进行字段重组：
    将除 'prompt' 和 'input' 字段之外的字段放入新字段 'ground_truth' 中。

    参数:
        file_path (str): JSONL 文件的路径。
    """
    if not isinstance(file_path, str) or not file_path.endswith('.jsonl'):
        raise ValueError("输入必须是一个有效的 .jsonl 文件路径")

    # 读取并处理文件内容
    updated_lines = []
    with open(file_path, 'r', encoding='utf-8') as file:
        for line in file:
            try:
                data = json.loads(line.strip())  # 解析每一行的 JSON 对象
                if not isinstance(data, dict):
                    raise ValueError("文件中的某一行不是有效的 JSON 对象")

                # 提取 'prompt' 和 'input' 字段（如果存在）
                result = {}
                if 'prompt' in data:
                    result['prompt'] = data['prompt']

                # 将其他字段放入 'ground_truth'
                ground_truth = {key: value for key, value in data.items() if key not in ('prompt')}
                if ground_truth:
                    result['ground_truth'] = ground_truth
                    result["data_source"] = 'Cipher'
                # 将重组后的 JSON 对象转换为字符串并保存
                updated_lines.append(json.dumps(result, ensure_ascii=False))
            except json.JSONDecodeError as e:
                raise ValueError(f"文件中的某一行无法解析为 JSON: {line.strip()}") from e
    random.shuffle(updated_lines)
    # 将更新后的内容写回文件
    with open(file_path, 'w', encoding='utf-8') as file:
        file.write('\n'.join(updated_lines))


def main_pipeline(cipher_envs, input_set, timestamp):

    for cipher_env in tqdm(cipher_envs, desc='cipher data generating'):
        generator = CipherDataGenerator(env_class=cipher_env,input_list=input_set)
        generator.write_to_jsonl(timestamp)

    final_data_path = fr'examples/bootcamp_generator_outputs/{timestamp}/cipher/final_data/'
    if not os.path.exists(final_data_path):
        os.makedirs(final_data_path)
    else:
        # print(f'Warning: {final_data_path} already exists')
        pass
    translate(fr'examples/bootcamp_generator_outputs/{timestamp}/cipher/origin_tmp_data', final_data_path)
    final_data_enhanced_path = final_data_path +'_enhanced'
    enhance(refined_rules = 'internbootcamp/libs/data/cipher_rules.jsonl',inputDir = final_data_path,outputDir = final_data_enhanced_path)
    from internbootcamp.bootcamp_utils import merge_jsonlines
    merge_output_file=fr"examples/bootcamp_generator_outputs/{timestamp}/cipher.jsonl"
    merge_jsonlines(input_directory=final_data_enhanced_path,output_file=merge_output_file)
    restructure_jsonl_file(merge_output_file)
    try:
        shutil.rmtree(fr'examples/bootcamp_generator_outputs/{timestamp}/cipher')
    except:
        pass
    return final_data_enhanced_path

if __name__ == '__main__':
    # 指定cipher环境列表
    # cipher_env_list为所有cipher环境
    # kor_cipher_env_list仅含kor cipher环境
    parser = argparse.ArgumentParser(description='Process...')
    parser.add_argument('--nums', type=int)
    parser.add_argument('--split', type=str)
    parser.add_argument('--timestamp', type=str)
    parser.add_argument('--filepath', type=str)
    parser.add_argument('--coef', type=float, default=30/144, help='Coef for balance cipher data tasks')
    args = parser.parse_args()

    # 获取传入的参数
    nums = int(args.nums)
    if nums == 0:
        exit()
    nums = int(nums * args.coef)
    split = str(args.split)
    timestamp = str(args.timestamp)
    # df = pd.read_csv(r'internbootcamp/libs/data/train.tsv', sep='\t')
    # df = pd.read_csv(args.filepath, sep='\t')
    sentences = open(args.filepath, 'r', encoding='utf-8').readlines()
    
    cipher_env_list.remove(DoubleCeaserEnvironment)
    cipher_env_list.remove(KorPigpenMasonicCipherEnvironment)
    main_pipeline(cipher_envs=cipher_env_list, input_set = sspd(sentences, percentage_distribution={(2,20):1},total_samples=nums), timestamp=timestamp + f"/{split}")