datasets.py

import json
import random
import logging
import pandas as pd
import os
from tqdm import tqdm
class CodeBlock(object):
    def __init__(self, file_path, description, code_content, language, _type, _api_block=False):
        """
        Represents a block of code.
        :param file_path: The path to the code file.
        :param description: The description of the code block.
        :param code_content: The content of the code block.
        """
        self.file_path = file_path
        self.code_content = code_content
        self.description = description
        self.language = language
        self._type = _type
        self._api_block = _api_block

    def __str__(self):
        if self.language == "python":
            comment_label = "#"
        else:
            comment_label = "//"
        
        crossfile_context = "\n".join([f"{comment_label} {cl}" for cl in self.description.strip().split('\n') if cl]) + "\n"
        crossfile_context += "\n".join([f"{comment_label} {cl}" for cl in  self.code_content.split('\n') if cl])
        return crossfile_context.strip()

class Example(object):
    def __init__(self, task_id, file_path, left_context, right_context, related_files, target_code, language):
        """
        Represents an example used for constructing a dataset.
        :param task_id: Task ID.
        :param file_path: File path.
        :param left_context: The context to the left of the target code.
        :param right_context: The context to the right of the target code.
        :param related_files: A list of related files, each containing a path and text.
        :param target_code: The target code snippet.
        """
        self.task_id = task_id
        self.file_path = file_path
        self.left_context = left_context
        self.right_context = right_context
        self.related_files = related_files
        self.target_code = target_code
        self.language = language

    def __str__(self):
        return f"[Example]:\n[Task ID]:\n{self.task_id}\n[Path]:\n{self.file_path}\n[Left Context]:\n{self.left_context}\n[Target Code]:\n{self.target_code}\n[Right Context]:\n{self.right_context}\n[Related Files]:\n{len(self.related_files)} files\n"

    
def load_test_dataset(args, dataset_name, language):
    """
    Loads a dataset.
    :param args: Parameters containing various configurations.
    :param datasetname: The name of the dataset to load.
    :param language: The language of the data to load.
    :return: The loaded dataset.
    """
    if dataset_name =="cceval":
        data_frame = pd.read_parquet(f"~/research/Data4AlignCoder/data/{dataset_name}/{language}/test.parquet")
        # data_frame['crossfile_context'] = data_frame['crossfile_context'].apply(json.loads)
        # data_frame['crossfile_context'] = []
    elif dataset_name == 'repoeval_update':
        data_frame = pd.read_parquet(f"data/{dataset_name}/{language}/python/test.parquet")
        data_frame['crossfile_context'] = data_frame['crossfile_context'].apply(json.loads)
    elif dataset_name == 'repoeval' and language != 'func_level':
        data_frame1 = pd.read_parquet(f"data/{dataset_name}/{language}/test_0.parquet")
        data_frame2 = pd.read_parquet(f"data/{dataset_name}/{language}/test_1.parquet")
        data_frame = pd.concat([data_frame1, data_frame2], ignore_index=True)
    else:
        data_frame = pd.read_parquet(f"data/{dataset_name}/{language}/test.parquet")

    if dataset_name == 'repoeval':
        language = 'python'

    if args.debug:
        data_frame = data_frame.sample(1000)
    
    dataset = []
    for item in data_frame[["task_id", "path", "left_context", "right_context", "crossfile_context", "groundtruth"]].values:
        cross_files = item[4] if len(item[4]) > 0 else [{'path': "", "text": "Don't need cross file context for completion"}]
        cross_files = [CodeBlock(x["path"], f"file path: {x['path']}\nlines: {0}-{len(x['text'].splitlines())}", x["text"], language, '') for x in cross_files]
        dataset.append(Example(item[0], item[1], item[2], item[3], cross_files, item[5], language))
    
    return dataset
def load_harmony_test_dataset(file_path):
    """
    load harmony test dataset
    """
    dataset_path = [
        "BooksAndReferenceTemplate", "BusinessTemplate", "CarsTemplate", 
        "EducationTemplate", "EntertainmentTemplate", "FinanceTemplate", 
        "FoodAndDrinkTemplate", "HealthAndFitnessTemplate", "HouseAndHomeTemplate", 
        "KidsTemplate", "LifestyleAndServiceTemplate", "MedicalTemplate", 
        "MovieTVAndLivestreamingTemplate", "NavigationTemplate", "NewsTemplate", 
        "PhotographyTemplate", "ShoppingTemplate", "SocialTemplate", 
        "ToolsTemplate", "TravelTemplate"
    ]
    dataset = []
    for file in dataset_path:
        count = 0
        with open(f"{file_path}/{file}_api_code_data.json", "r") as f:
            data = json.load(f)
            # 修复：data是字典，需要遍历其值而不是键
            for data_key, item in data.items():
                cross_files = item["related_files"] if len(item["related_files"]) > 0 else [{'file_path': "", "code_content": "Don't need cross file context for completion"}]
                cross_files = [CodeBlock(x["file_path"], f"file path: {x['file_path']}\nlines: {0}-{len(x['code_content'].splitlines())}", x["code_content"], "arkts", '') for x in cross_files]
                # 使用 template + task_id 作为唯一标识
                dataset.append(Example(file + item["task_id"], item["file_path"], item["left_context"], item["right_context"], cross_files, item.get("target_code", ""), "arkts"))

    return dataset

def load_train_and_valid_dataset(data_root: str | None = None):
    """
    Loads the training dataset.
    :return: The training dataset.
    """
    if data_root is None:
        data_root = os.environ.get("ALIGNCODER_DATA_ROOT")
    if not data_root:
        raise ValueError(
            "Missing dataset root. Pass `data_root` or set env var `ALIGNCODER_DATA_ROOT` "
            "to the directory that contains `data/github_repos/...`."
        )

    training_datasets, validation_datasets = [], []
    validation_datasets = []
    for language in ["python", "java"]:
        data_frame = pd.read_parquet(f"{data_root}/data/github_repos/{language}/train.parquet")
        all_data, temp_data = [], []
        for x in data_frame[["path", "content", "first"]].values:
            if x[-1]:  # At the start of a new file
                if len(temp_data) > 1:
                    all_data.append((temp_data, language))
                temp_data = []
            temp_data.append([x[0], x[1]])
        training_datasets.extend(all_data[:2000])
        validation_datasets.extend(all_data[2000:2200])
    
    random.shuffle(training_datasets)
    random.shuffle(validation_datasets)

    return training_datasets, validation_datasets
def load_train_and_valid_harmony_dataset(path):
    """
    load harmony dataset
    """
    training_datasets, validation_datasets = [], []
    files = sorted(os.listdir(path))
    count = 0
    for file in tqdm(files):
        with open(os.path.join(path, file), "r") as f:
            data = json.load(f)
            if count < 16:
                training_datasets.append(data)
            else:
                validation_datasets.append(data)
        count += 1
    real_training_datasets, real_validation_datasets = [], []
    for datas in training_datasets:
        for data in datas.values():
            real_training_datasets.append(data)
    for datas in validation_datasets:
        for data in datas.values():
            real_validation_datasets.append(data)
    random.shuffle(real_training_datasets)
    random.shuffle(real_validation_datasets)
    return real_training_datasets, real_validation_datasets

def construct_dataset(raw_data, num_samples):
    """
    Builds a dataset.
    :param raw_data: Raw data.
    :param num_samples: The number of samples to generate.
    :return: The list of constructed samples.
    """
    examples = []
    data_index = 0
    while len(examples) < num_samples:
        example, language = raw_data[data_index % len(raw_data)]
        data_index += 1
        selected_file = random.choice(example[1:])
        related_files = [CodeBlock(x[0], f"file path: {x[0]}\nlines: {0}-{len(x[1].splitlines())}", x[1], language, '') for x in example if x[0] != selected_file[0]]
        path = selected_file[0]
        selected_file_content = selected_file[1].split(" ")
        try_count = 0
        while try_count < 10:
            end_line_number = int(len(selected_file_content) * random.uniform(0.1, 0.8))
            left_context = " ".join(selected_file_content[:end_line_number])
            target_length = random.randint(16, 96)
            target = " ".join(selected_file_content[end_line_number:end_line_number + target_length])
            right_context = " ".join(selected_file_content[end_line_number + target_length:])
            if len(left_context.split()) > 30 and len(target.split()) > 8:
                examples.append(Example(len(examples), path, left_context, right_context, related_files, target, language))
                break
            try_count += 1
    print("data index", data_index)
    return examples

def construct_harmony_dataset(raw_data, num_samples,rand = 0):
    """
    Builds a dataset.
    :param raw_data: Raw data.
    :param num_samples: The number of samples to generate.
    :return: The list of constructed samples.
    """
    examples = [] 
    data_index = 0
    while len(examples) < num_samples:
        if rand == 0:
            example = raw_data[data_index % len(raw_data)]
        else:
            example = random.choice(raw_data)
        # example = raw_data[data_index % len(raw_data)]
        language = example["language"]
        data_index += 1
        files = []

        files.append({
            "file_path": example["file_path"],
            "code_content": example["code"],
            "language": language
        })

        for file in example["related_files"]:
            files.append({
                "file_path": file["file_path"],
                "code_content": file["code_content"],
                "language": language
            })
        if len(files) == 0:
            continue
        selected_file = random.choice(files)
        related_files = [CodeBlock(x["file_path"], f"file path: {x['file_path']}\nlines: {0}-{len(x['code_content'].splitlines())}", x["code_content"], language, '') for x in files if x["file_path"] != selected_file["file_path"]]
        path = selected_file["file_path"]
        selected_file_content = selected_file["code_content"].split(" ")
        try_count = 0
        while try_count < 10:
            end_line_number = int(len(selected_file_content) * random.uniform(0.1, 0.8))
            left_context = " ".join(selected_file_content[:end_line_number])
            target_length = random.randint(16, 96)
            target = " ".join(selected_file_content[end_line_number:end_line_number + target_length])
            right_context = " ".join(selected_file_content[end_line_number + target_length:])
            if len(left_context.split()) > 30 and len(target.split()) > 8:
                examples.append(Example(len(examples), path, left_context, right_context, related_files, target, language))
                break
            try_count += 1
    # print("data index", data_index)        
    return examples

if __name__ == "__main__":
    dataset = load_harmony_test_dataset()
    print(dataset)