import datasets 
import fire 

from common .common_types import Sample 
from common .dataset import create_datasets 


def load_livecodebench ()->list [Sample ]:

    rows =datasets .load_dataset ("livecodebench/execution-v2",split ="test")
    samples =[]
    for row in rows :
        fn_name =row ["function_name"]
        call_str =row ["input"]


        assert call_str .startswith (fn_name +"(")
        assert call_str .endswith (")")
        input_str =call_str [len (fn_name )+1 :-1 ]

        samples .append (
        Sample (
        sample_id =str (row ["id"]),
        code =row ["code"],
        input =input_str ,
        output =row ["output"],
        function_name =fn_name ,
        )
        )
    return samples 


def main (data_name :str ="livecodebench",output_dir :str ="/path/to/home/lltm/02_codeexec_etcot/scripts/instruction/convert_datasets"):
    all_samples =load_livecodebench ()
    create_datasets (data_name ,all_samples ,train_ratio =0 ,output_dir =output_dir )


if __name__ =="__main__":
    fire .Fire (main )
