chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,368 @@
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import gc
|
||||
import torch
|
||||
import argparse
|
||||
import random
|
||||
from hashlib import md5
|
||||
import multiprocessing as mp
|
||||
from typing import List, Optional
|
||||
|
||||
from constant import TaskType, Language, CodeLanguage, NUM_HARD_NEGATIVES
|
||||
from corpus_generator import CorpusGenerator
|
||||
from triplet_generator import TripletGenerator
|
||||
from search import get_top1
|
||||
|
||||
|
||||
def compute_md5(text: str):
|
||||
return md5(text.encode()).hexdigest()
|
||||
|
||||
|
||||
def get_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
'--task_type',
|
||||
type=str,
|
||||
required=True,
|
||||
help='The task type to generate data for',
|
||||
choices=[t.name for t in TaskType]
|
||||
)
|
||||
parser.add_argument(
|
||||
'--code_language',
|
||||
type=str,
|
||||
required=True,
|
||||
help='The code language to generate questions for.',
|
||||
choices=[c.name for c in CodeLanguage]
|
||||
)
|
||||
parser.add_argument(
|
||||
'--corpus_root',
|
||||
type=str,
|
||||
required=True,
|
||||
help='The root directory of the corpus data.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--save_dir',
|
||||
type=str,
|
||||
required=True,
|
||||
help='The path to save the generated data'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--examples_dir',
|
||||
type=str,
|
||||
default=None,
|
||||
help='The path to the examples directory. If not None, the examples will be used for few-shot generation.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--num_examples',
|
||||
type=int,
|
||||
default=3,
|
||||
help='The number of examples to use for few-shot generation. Default: 3'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--cache_dir',
|
||||
type=str,
|
||||
default=None,
|
||||
help='The cache directory'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--language',
|
||||
type=str,
|
||||
default='en',
|
||||
help='The language to generate for. ISO 639-1 code. Default: en',
|
||||
choices=[l.name for l in Language]
|
||||
)
|
||||
parser.add_argument(
|
||||
'--tgt_code_language',
|
||||
type=str,
|
||||
default=None,
|
||||
help='The target code language to generate code translations for.',
|
||||
choices=[c.name for c in CodeLanguage]
|
||||
)
|
||||
parser.add_argument(
|
||||
'--num_samples',
|
||||
type=int,
|
||||
default=-1,
|
||||
help='The number of examples to use for generation. Default: -1. Use all available examples.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--model',
|
||||
type=str,
|
||||
default='Qwen2.5-72B-Instruct',
|
||||
help='The model to use for generation. Default: Qwen2.5-72B-Instruct'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--model_type',
|
||||
type=str,
|
||||
default='open-source',
|
||||
help='The type of model to use for generation. Default: open-source',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--port',
|
||||
type=int,
|
||||
default=8000,
|
||||
help='The port for vllm.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--num_processes',
|
||||
type=int,
|
||||
default=1,
|
||||
help='The number of processes to use for generation. Default: 1'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--doc_length',
|
||||
type=str,
|
||||
default='len_0_500',
|
||||
help='The corpus length used to load dataset. Default: len_0_500'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--external_path',
|
||||
type=str,
|
||||
default='',
|
||||
help='The corpus length used to load dataset. Default: len_0_500'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--sim_model_name',
|
||||
type=str,
|
||||
default=None,
|
||||
help='The language of source corpus.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--max_corpus',
|
||||
type=int,
|
||||
default=500000,
|
||||
help='The max num of corpus to load.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--overwrite',
|
||||
action='store_true',
|
||||
help='Whether to overwrite the existing data.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--debug_mode',
|
||||
action='store_true',
|
||||
help='Whether to open debug mode.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--gen_hard_neg',
|
||||
action='store_true',
|
||||
help='Whether to generate hard negatives.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'--seed',
|
||||
type=int,
|
||||
default=None,
|
||||
help='Random seed for generating triplets using the same positive. Default: 42'
|
||||
)
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def gen_triplets(
|
||||
model: str,
|
||||
model_type: str,
|
||||
port: int,
|
||||
positives: List[dict],
|
||||
task_type: str,
|
||||
language: str,
|
||||
code_language: str,
|
||||
tgt_code_language: str,
|
||||
examples_pool: Optional[List[dict]] = None,
|
||||
num_examples: int = 3,
|
||||
tqdm_desc: str = "Generating triplets",
|
||||
thread_count: int = 1,
|
||||
gen_cache_dir: Optional[str] = None,
|
||||
debug_mode: bool = False,
|
||||
gen_hard_neg: bool = False,
|
||||
):
|
||||
triplet_generator = TripletGenerator(model, model_type, port, cache_dir=gen_cache_dir)
|
||||
triplets = triplet_generator.run(
|
||||
positives=positives,
|
||||
task_type=task_type,
|
||||
language=language,
|
||||
code_language=code_language,
|
||||
tgt_code_language=tgt_code_language,
|
||||
examples_pool=examples_pool,
|
||||
num_examples=num_examples,
|
||||
tqdm_desc=tqdm_desc,
|
||||
thread_count=thread_count,
|
||||
debug_mode=debug_mode,
|
||||
gen_hard_neg=gen_hard_neg,
|
||||
num_negatives=NUM_HARD_NEGATIVES,
|
||||
)
|
||||
return triplets
|
||||
|
||||
|
||||
def get_save_path(
|
||||
save_dir: str,
|
||||
task_type: str,
|
||||
language: str,
|
||||
code_language: str,
|
||||
tgt_code_language: Optional[str] = None
|
||||
):
|
||||
save_dir = os.path.join(save_dir, language, task_type)
|
||||
if tgt_code_language is not None:
|
||||
file_name = f"{language}-{code_language}-to-{tgt_code_language}-triplets.jsonl"
|
||||
else:
|
||||
file_name = f"{language}-{code_language}-triplets.jsonl"
|
||||
save_path = os.path.join(save_dir, file_name)
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
return save_path
|
||||
|
||||
|
||||
def save_triplets(
|
||||
triplets: list,
|
||||
save_dir: str,
|
||||
task_type: str,
|
||||
language: str,
|
||||
code_language: str,
|
||||
tgt_code_language: Optional[str] = None
|
||||
):
|
||||
if len(triplets) == 0:
|
||||
print(f"No triplets to save: {task_type} | {language} | {code_language} | {tgt_code_language}")
|
||||
return
|
||||
|
||||
save_path = get_save_path(save_dir, task_type, language, code_language, tgt_code_language)
|
||||
query_md5s = set()
|
||||
pos_md5s = set()
|
||||
old_triplets = []
|
||||
if os.path.exists(save_path):
|
||||
with open(save_path, "r", encoding="utf-8") as f:
|
||||
for line in f.readlines():
|
||||
triplet = json.loads(line)
|
||||
old_triplets.append(triplet)
|
||||
query_md5s.add(compute_md5(triplet['query']))
|
||||
pos_md5s.add(compute_md5(triplet['pos'][0]))
|
||||
|
||||
with open(save_path, 'w', encoding='utf-8') as f:
|
||||
for triplet in old_triplets:
|
||||
f.write(json.dumps(triplet, ensure_ascii=False) + '\n')
|
||||
|
||||
for triplet in triplets:
|
||||
_query_md5 = compute_md5(triplet['query'])
|
||||
_pos_md5 = compute_md5(triplet['pos'][0])
|
||||
if _query_md5 in query_md5s or _pos_md5 in pos_md5s:
|
||||
continue
|
||||
f.write(json.dumps(triplet, ensure_ascii=False) + '\n')
|
||||
print(f"Triplets saved to {save_path}")
|
||||
|
||||
|
||||
def main(args):
|
||||
# set seed
|
||||
seed = args.seed
|
||||
if seed is not None:
|
||||
print(f"------------------- Seed set to {seed} -------------------")
|
||||
random.seed(seed)
|
||||
|
||||
model = args.model
|
||||
model_type = args.model_type
|
||||
port = args.port
|
||||
|
||||
num_samples = args.num_samples
|
||||
|
||||
task_type = args.task_type
|
||||
language = args.language
|
||||
code_language = args.code_language
|
||||
tgt_code_language = args.tgt_code_language
|
||||
|
||||
corpus_root = args.corpus_root
|
||||
corpus_dir = os.path.join(corpus_root, code_language)
|
||||
doc_length = args.doc_length.split()
|
||||
external_path = args.external_path.split()
|
||||
|
||||
save_dir = args.save_dir
|
||||
cache_dir = args.cache_dir
|
||||
num_processes = min(args.num_processes, int(mp.cpu_count() * 0.8))
|
||||
overwrite = args.overwrite
|
||||
debug_mode = args.debug_mode
|
||||
gen_hard_neg = args.gen_hard_neg
|
||||
|
||||
save_path = get_save_path(save_dir, task_type, language, code_language, tgt_code_language)
|
||||
# if os.path.exists(save_path) and not overwrite:
|
||||
# data = []
|
||||
# with open(save_path) as f:
|
||||
# for line in f:
|
||||
# data.append(json.loads(line))
|
||||
# if len(data) >= num_samples * 0.8:
|
||||
# print(f"Triplets already exist at {save_path}. Skipping generation.")
|
||||
# return
|
||||
# else:
|
||||
# print(f"Triplets already exist at {save_path}. But samples is really small, continue generation.")
|
||||
# num_samples = int((num_samples - len(data)) * 1.25) # consider the filtered samples
|
||||
|
||||
corpus_generator = CorpusGenerator(cache_dir)
|
||||
|
||||
examples_dir = args.examples_dir
|
||||
num_examples = args.num_examples
|
||||
if examples_dir is not None:
|
||||
# if task_type in ["single_turn_code_qa", "multi_turn_code_qa"]:
|
||||
# examples_path = os.path.join(examples_dir, language, task_type, "sample_examples.json")
|
||||
if task_type in ["code_translation_retrieval"]:
|
||||
examples_path = os.path.join(examples_dir, language, task_type,
|
||||
f"{code_language}-to-{tgt_code_language}_sample_examples.json")
|
||||
else:
|
||||
examples_path = os.path.join(examples_dir, language, task_type, f"{code_language}_sample_examples.json")
|
||||
try:
|
||||
with open(examples_path, 'r', encoding='utf-8') as f:
|
||||
examples_pool = json.load(f)
|
||||
examples_pool = random.sample(examples_pool,
|
||||
min(30, len(examples_pool))) # sample 30 examples for few-shot generation
|
||||
except:
|
||||
print(f'Error for loading examples from {examples_path}')
|
||||
examples_pool = None
|
||||
else:
|
||||
examples_pool = None
|
||||
|
||||
positives, large_positives = corpus_generator.run(
|
||||
num_samples=num_samples,
|
||||
max_corpus=args.max_corpus,
|
||||
corpus_dir=corpus_dir,
|
||||
doc_length=doc_length,
|
||||
external_path=external_path,
|
||||
source_language=code_language
|
||||
)
|
||||
|
||||
if task_type in ["code_modification_retrieval", "code_comparison_retrieval"]:
|
||||
top1_docs = get_top1([e['text'] for e in positives], args.sim_model_name, [e['text'] for e in large_positives])
|
||||
for i in range(len(top1_docs)):
|
||||
positives[i]['similar'] = top1_docs[i]
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
print("=================== Generate training data ===================")
|
||||
print(f'Task Type: {task_type} | Language: {language} | Code Language: {code_language} | Target Code Language: {tgt_code_language}')
|
||||
start_time = time.time()
|
||||
triplets = gen_triplets(
|
||||
model=model,
|
||||
model_type=model_type,
|
||||
port=port,
|
||||
positives=positives,
|
||||
task_type=task_type,
|
||||
language=language,
|
||||
code_language=code_language,
|
||||
tgt_code_language=tgt_code_language,
|
||||
examples_pool=examples_pool,
|
||||
num_examples=num_examples,
|
||||
thread_count=num_processes,
|
||||
gen_cache_dir=os.path.join(save_dir, language, task_type, "gen_cache_dir"),
|
||||
debug_mode=debug_mode,
|
||||
gen_hard_neg=gen_hard_neg,
|
||||
)
|
||||
save_triplets(
|
||||
triplets=triplets,
|
||||
save_dir=save_dir,
|
||||
task_type=task_type,
|
||||
language=language,
|
||||
code_language=code_language,
|
||||
tgt_code_language=tgt_code_language
|
||||
)
|
||||
end_time = time.time()
|
||||
print("=============================================================")
|
||||
print(f"Time taken: {end_time - start_time:.2f} seconds")
|
||||
print("=============================================================")
|
||||
print("DONE!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = get_args()
|
||||
main(args)
|
||||
Reference in New Issue
Block a user