generation.py

from utils import *

def generate_output(model_folder, api_model, carbon_intensity, tokenizer, Generation_Agent_JSON, Generation_Agent_TOON, Generation_Agent_XML, Generation_Agent_YAML, no_system_prompt=False):
    formats = ["JSON", "YAML", "XML", "TOON(JSON)", "TOON(YAML)", "TOON(XML)"]

    for selected_format in formats:
        selected_path = f"task/Text_to_{selected_format}.json"
        output_dir_path = f"output/{model_folder}/{selected_format}"

        if "TOON" in selected_format.upper():
            system_prompt_path = f"prompts/system_prompt_TOON.txt"
        else:
            system_prompt_path = f"prompts/system_prompt_{selected_format}.txt"

        with open(system_prompt_path, "r", encoding="utf-8") as f:
            system_prompt_txt = f.read()

        input_tokens = tokenizer.encode(system_prompt_txt, add_special_tokens=False)
        input_num_tokens = len(input_tokens)

        os.makedirs(output_dir_path, exist_ok=True)
        os.makedirs(f"{output_dir_path}/raw", exist_ok=True)

        data = load_json(selected_path)
        
        results_csv_path = f"{output_dir_path}/results.csv"
        gcs_csv_path = f"{output_dir_path}/GCS_env_df"
        processed_ids = set()
        
        if os.path.exists(results_csv_path):
            existing_df = pd.read_csv(results_csv_path, dtype={"ID": str})
            processed_ids = set(existing_df["ID"])
            result_df = existing_df
        else:
            result_df = pd.DataFrame(columns=["ID", "CO2(Orig)", "Energy_Consumption", "NT(Input)", "Duration", "NT", "CE", "Render_Score", "Syntax_Score", "GCS", "EES"])
        
        if os.path.exists(gcs_csv_path):
            gcs_df = pd.read_csv(gcs_csv_path)
        else:
            gcs_df = pd.DataFrame(columns=BASE_COLUMNS + GCS_ENV_COLUMNS)

        data_to_process = [item for item in data if str(item.get("task_id")) not in processed_ids]

        for item in tqdm(data_to_process, desc=f"{model_folder} - {selected_format}"):
            task_id = item.get("task_id")

            query = item.get("query")
            raw_output_metric = item.get("raw_output_metric", [])
            generated_code_file_path = f"{output_dir_path}/raw/{task_id}.txt"

            render_score = 0
            syntax_score = 0
            final_score = 0
            parsed_success = 0
            answer = ""
            generated_keys = []

            start_time = time.time()

            if no_system_prompt:
                query_with_system = f"{system_prompt_txt}\n\n{query}"
            else:
                query_with_system = query

            if "TOON" in selected_format.upper():
                user_proxy.initiate_chat(Generation_Agent_TOON, message=query_with_system)
                answer = user_proxy.last_message(Generation_Agent_TOON)["content"]
            elif selected_format == "JSON":
                user_proxy.initiate_chat(Generation_Agent_JSON, message=query_with_system)
                answer = user_proxy.last_message(Generation_Agent_JSON)["content"]
            elif selected_format == "XML":
                user_proxy.initiate_chat(Generation_Agent_XML, message=query_with_system)
                answer = user_proxy.last_message(Generation_Agent_XML)["content"]
            elif selected_format == "YAML":
                user_proxy.initiate_chat(Generation_Agent_YAML, message=query_with_system)
                answer = user_proxy.last_message(Generation_Agent_YAML)["content"]
            
            answer = clean_llm_response(answer)

            end_time = time.time()
            duration = end_time - start_time

            tokens = tokenizer.encode(answer, add_special_tokens=False)
            output_num_tokens = len(tokens)
            
            emissions_data = calculate_api_emissions(
                model_name=api_model,
                input_tokens=input_num_tokens,
                output_tokens=output_num_tokens,
                duration=duration,
                carbon_intensity=carbon_intensity
            )

            if os.path.exists(generated_code_file_path):
                os.remove(generated_code_file_path)

            with open(generated_code_file_path, "w", encoding="utf-8") as f:
                f.write(answer)

            if selected_format == "JSON":
                try:
                    with open(generated_code_file_path, 'r', encoding='utf-8') as f:
                        result = json.load(f)
                        generated_keys = extract_keys_json(result)
                    parsed_success = 1
                except Exception as e:
                    parsed_success = 0
            elif selected_format == "XML":
                try:
                    with open(generated_code_file_path, 'r', encoding='utf-8') as f:
                        result = ET.fromstring(f.read())
                        generated_keys = extract_keys_xml(result)
                    parsed_success = 1
                except Exception as e:
                    parsed_success = 0
            elif selected_format == "YAML":
                try:
                    with open(generated_code_file_path, 'r', encoding='utf-8') as f:
                        result = yaml.safe_load(f)
                        generated_keys = extract_keys_yaml(result)
                    parsed_success = 1
                except Exception as e:
                    parsed_success = 0
            elif "TOON" in selected_format.upper():
                try:
                    result = decode(answer)
                    generated_keys = extract_keys_json(result)
                    parsed_success = 1
                except ToonDecodeError as e:
                    parsed_success = 0

            if parsed_success != 0:
                render_score = 1
                intersection = list(set(generated_keys) & set(raw_output_metric))
                syntax_score = len(intersection) / len(raw_output_metric)

                final_score = (0.2 * render_score) + (0.8 * syntax_score)

            CO2 = (emissions_data["emissions"] * output_num_tokens) / (input_num_tokens + output_num_tokens)
            X = (1000 * CO2) / output_num_tokens
            E_env = 1 / (1 + (X / 0.001719))

            result_df.loc[len(result_df)] = {
                "ID": task_id,
                "CO2(Orig)": round(emissions_data["emissions"], 8),
                "Energy_Consumption": round(emissions_data["energy_consumed"], 8),
                "NT(Input)": input_num_tokens,
                "Duration": round(duration, 2),
                "NT": output_num_tokens,
                "CE": round(CO2, 8),
                "Render_Score": round(render_score, 2),
                "Syntax_Score": round(syntax_score, 2),
                "GCS": round(final_score, 2),
                "EES": round(E_env, 2),
            }

            element_row = {"ID": task_id}

            for d in D_VALUES:
                element_row[f"GCS_env({d:.2f})"] = round((1 - d) * final_score + d * E_env, 2)

            gcs_df.loc[len(gcs_df)] = element_row

            result_df.to_csv(f"{output_dir_path}/results.csv", index=False)
            gcs_df.to_csv(f"{output_dir_path}/GCS_env_df", index=False)