U g$:e@s*ddlZddlZddlZGdddZdS)Nc@s0eZdZd ddZdddZdd Zd d Zd S)HumanEvalDatasetpythonFcCs||_ttj|jd|d|_||j|}g|_t t |D]"}t |D]}|j ||qTqH|jdd|_ t jdtd|dt |jdS) z root: the path to the HumanEval dataset sample_num: the number of samples for each prompt language: the language of the HumanEval dataset issft: whether to use the SFT setting z humaneval-z.jsonlr stopwordsizRead HumanEval from z, number of samples N)rootopenospathjoin readlinesdataget_qa_only_data clean_datarangelenappendrnprandomseedprint)selfr sample_numlanguageissfttmpijr^/weka-jd/prod/containers/zhuqihao/dev-cpu/deepseek-coder/Evaluation/HumanEval/utils/dataset.py__init__s"  zHumanEvalDataset.__init__cCsg}|D]r}t|}|d}d|kr4|d}n|d}|rLd|d}d|kr^|d}ng}|||d||dq|S)z data_json: the jsonl file of HumanEval sft: whether to use the SFT setting return: a list of dict, each dict contains the prompt, task_id and stopwords promptprefixzBelow is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request. ### Instruction: Write a program to perform the given task. Input: z ### Response: Z stop_tokenstask_id)r r"original_promptr)jsonloadsstripr)rZ data_jsonsftansliner Z origin_promptsrrrr s     z!HumanEvalDataset.get_qa_only_datacCs t|jS)z= return the number of samples in the dataset )rr)rrrr__len__2szHumanEvalDataset.__len__cCs|j|}|S)z, return the sample at index )r)rindexsamplerrr __getitem__8s zHumanEvalDataset.__getitem__N)rrF)F)__name__ __module__ __qualname__rr r+r.rrrrrs  r)rnumpyrr$rrrrrs