NexQuant/rdagent/components/coder/finetune/__init__.py

"""
LLM Fine-tuning CoSTEER Implementation

This module provides fine-tuning specific components for the CoSTEER framework,
including evaluators and evolving strategies.
"""

import json
from pathlib import Path
from typing import Callable

import yaml

from rdagent.app.finetune.llm.conf import FT_RD_SETTING
from rdagent.components.coder.CoSTEER import CoSTEER
from rdagent.components.coder.CoSTEER.evaluators import (
    CoSTEERMultiEvaluator,
    CoSTEERSingleFeedback,
)
from rdagent.components.coder.CoSTEER.evolving_strategy import (
    MultiProcessEvolvingStrategy,
)
from rdagent.components.coder.CoSTEER.knowledge_management import (
    CoSTEERQueriedKnowledge,
)
from rdagent.components.coder.finetune.conf import (
    FT_DATA_SCRIPT_NAME,
    FT_PATHS,
    FT_TEST_PARAMS_FILE_NAME,
    FT_YAML_FILE_NAME,
    FTCoderCoSTEERSettings,
)
from rdagent.components.coder.finetune.eval import FTCoderEvaluator, FTDataEvaluator
from rdagent.core.experiment import FBWorkspace, Task
from rdagent.core.scenario import Scenario
from rdagent.log import rdagent_logger as logger
from rdagent.oai.llm_utils import APIBackend
from rdagent.scenarios.finetune.scen.llama_factory_manager import LLaMAFactory_manager
from rdagent.scenarios.finetune.scen.utils import FinetuneDatasetDescriptor
from rdagent.utils.agent.tpl import T

DIRNAME = Path(__file__).absolute().resolve().parent


class LLMFinetuneEvolvingStrategy(MultiProcessEvolvingStrategy):
    """LLM Fine-tuning specific evolving strategy"""

    def __init__(self, scen: Scenario, settings, *args, **kwargs):
        super().__init__(scen, settings)
        self.llama_factory_manager = LLaMAFactory_manager

    def implement_func_list(self) -> list[Callable]:
        return [self.implement_data, self.implement_lf_config]

    def implement_data(
        self,
        target_task: Task,
        queried_knowledge: CoSTEERQueriedKnowledge | None = None,
        workspace: FBWorkspace | None = None,
        prev_task_feedback: CoSTEERSingleFeedback | None = None,
    ) -> dict[str, str]:
        """Generate data processing script based on task.

        This method generates a Python script that processes seed datasets
        and outputs a data.json file in Alpaca format.

        Returns:
            dict with "process_data.py" key containing the script code,
            or empty dict if data already exists.
        """
        # Check if proposal decided to skip data processing (reuse SOTA's data processing script)
        if getattr(target_task, "skip_data_processing", False):
            # Defensive check: ensure data script actually exists before skipping
            script_exists = False
            if workspace is not None:
                script_exists = FT_DATA_SCRIPT_NAME in workspace.file_dict

            if script_exists:
                logger.info("Proposal decided to skip data processing, reusing SOTA's data script")
                return {}
            else:
                logger.warning(
                    "skip_data_processing=True but process_data.py not found in workspace, "
                    "this indicates SOTA injection failed - system design issue"
                )
                # Don't fallback silently, let it fail early to expose the issue

        # check whether the current code passes evaluation
        if (
            prev_task_feedback is not None
            and "FTDataEvaluator" in prev_task_feedback.source_feedback
            and prev_task_feedback.source_feedback["FTDataEvaluator"]
        ):
            logger.info("Previous data processing code passed evaluation, skipping regeneration")
            return {}

        # build former failed trace
        queried_former_failed_knowledge = (
            queried_knowledge.task_to_former_failed_traces[target_task.get_task_information()]
            if queried_knowledge is not None
            else []
        )
        queried_former_failed_knowledge = (
            [
                knowledge
                for knowledge in queried_former_failed_knowledge[0]
                if knowledge.implementation.file_dict.get(FT_YAML_FILE_NAME)
                != workspace.file_dict.get(FT_YAML_FILE_NAME)
            ],
            queried_former_failed_knowledge[1],
        )

        # Get dataset information for the task
        involving_datasets = getattr(target_task, "involving_datasets", [])
        dataset_info = self._get_dataset_info(involving_datasets, datasets_path=FT_PATHS.datasets)

        # Generate data processing script using LLM
        system_prompt = T(".prompts:data_coder.system").r(
            scenario=self.scen.get_scenario_all_desc(),
            task_desc=target_task.get_task_information(),
            dataset_info=dataset_info,
            queried_former_failed_knowledge=queried_former_failed_knowledge[0],
            api_max_workers=FT_RD_SETTING.api_max_workers,
            datasets_path=FT_PATHS.datasets,
            workspace_path=FT_PATHS.workspace,
            force_think_token=FT_RD_SETTING.force_think_token,
        )

        user_prompt = T(".prompts:data_coder.user").r(
            datasets_path=FT_PATHS.datasets,
            workspace_path=FT_PATHS.workspace,
            latest_code=workspace.file_dict.get(FT_DATA_SCRIPT_NAME, "") if workspace else "",
            latest_feedback=prev_task_feedback,
            involved_dataset_folder_desc={
                ds_name: FinetuneDatasetDescriptor().describe_dataset_folder(
                    Path(FT_RD_SETTING.file_path) / "datasets" / ds_name, include_dataset_readme=True
                )
                for ds_name in involving_datasets
            },
        )

        script_code = APIBackend().build_messages_and_create_chat_completion(
            user_prompt=user_prompt,
            system_prompt=system_prompt,
            json_mode=False,
            code_block_language="python",
            code_block_fallback=False,
        )
        logger.info(f"Generated data processing script ({len(script_code)} chars)")

        return {FT_DATA_SCRIPT_NAME: script_code}

    def _get_dataset_info(self, involving_datasets: list[str], datasets_path: str = None) -> str:
        """Read dataset_info.json and return information for specified datasets.

        Handles unified tasks structure:
        - readme: Dataset README content
        - file_tree: Directory structure
        - total_samples: Total sample count
        - tasks: Dict of task info (use "_root" for root-level data files)

        Args:
            involving_datasets: List of dataset names to include
            datasets_path: Base path for datasets (e.g., "/assets/datasets/")
        """
        datasets_dir = Path(FT_RD_SETTING.file_path) / "datasets"
        dataset_info_path = datasets_dir / "dataset_info.json"

        # Use provided path or get from config
        if datasets_path is None:
            datasets_path = FT_PATHS.datasets

        if not dataset_info_path.exists():
            logger.warning(f"dataset_info.json not found at {dataset_info_path}")
            return "No dataset information available."

        try:
            with open(dataset_info_path, "r", encoding="utf-8") as f:
                all_dataset_info = json.load(f)
        except Exception as e:
            logger.error(f"Failed to read dataset_info.json: {e}")
            return f"Error reading dataset info: {e}"

        # Filter to only involved datasets, or use all if none specified
        if involving_datasets:
            filtered_info = {name: info for name, info in all_dataset_info.items() if name in involving_datasets}
        else:
            filtered_info = all_dataset_info

        if not filtered_info:
            return "No matching datasets found in dataset_info.json."

        # Format dataset info for the prompt
        info_parts = []
        for name, info in filtered_info.items():
            info_text = f"### Dataset: {name}\n"
            # IMPORTANT: Tell LLM the full path to dataset directory
            dataset_full_path = f"{datasets_path}{name}/"
            info_text += f"- **Dataset path**: `{dataset_full_path}` (each dataset has its own subdirectory)\n"
            info_text += f"- Total samples: {info.get('total_samples', 'N/A')}\n"
            info_text += f"- Size: {info.get('total_size_mb', 'N/A')} MB\n"

            # File tree for understanding directory structure
            if info.get("file_tree"):
                file_tree = info["file_tree"]
                # Truncate if too long
                if len(file_tree) > 1000:
                    file_tree = file_tree[:1000] + "\n..."
                info_text += f"\n**File Structure** (relative to `{dataset_full_path}`):\n```\n{file_tree}\n```\n"

            # Handle unified tasks structure
            tasks = info.get("tasks", {})
            if tasks:
                info_text += "\n**Tasks:**\n"
                for task_name, task_info in tasks.items():
                    # "_root" indicates data files are in root directory
                    display_name = "(root)" if task_name == "_root" else task_name
                    info_text += f"\n#### {display_name}\n"
                    # Show full paths for data files
                    files = task_info.get("files", [])
                    info_text += f"- Files: {files}\n"
                    if files:
                        info_text += f"  - Full path example: `{dataset_full_path}{files[0]}`\n"
                    info_text += f"- Sample count: {task_info.get('sample_count', 'N/A')}\n"
                    if task_info.get("column_stats"):
                        # Show key token stats
                        stats_summary = []
                        for col, stats in task_info["column_stats"].items():
                            if stats.get("p50_tokens", 0) > 0:
                                stats_summary.append(f"{col}: p50={stats['p50_tokens']}, p99={stats['p99_tokens']}")
                        if stats_summary:
                            info_text += f"- Token stats: {'; '.join(stats_summary[:5])}\n"

            # README excerpt
            if info.get("readme"):
                readme = info["readme"]
                if len(readme) > 500:
                    readme = readme[:500] + "..."
                info_text += f"\n**README:**\n{readme}\n"

            info_parts.append(info_text)

        return "\n".join(info_parts)

    def implement_lf_config(
        self,
        target_task: Task,
        queried_knowledge: CoSTEERQueriedKnowledge | None = None,
        workspace: FBWorkspace | None = None,
        prev_task_feedback: CoSTEERSingleFeedback | None = None,
    ) -> dict[str, str]:
        """Implement a single fine-tuning task by generating LlamaFactory config"""
        if prev_task_feedback is not None and prev_task_feedback.source_feedback.get("FTCoderEvaluator", False):
            logger.info("Previous training code passed evaluation, skipping regeneration")
            return {}

        task_info = target_task.get_task_information()

        queried_former_failed_knowledge = (
            queried_knowledge.task_to_former_failed_traces[task_info] if queried_knowledge is not None else []
        )
        queried_former_failed_knowledge = (
            [
                knowledge
                for knowledge in queried_former_failed_knowledge[0]
                if knowledge.implementation.file_dict.get(FT_YAML_FILE_NAME)
                != workspace.file_dict.get(FT_YAML_FILE_NAME)
            ],
            queried_former_failed_knowledge[1],
        )

        # Get task parameters from the task object
        base_model = getattr(target_task, "base_model")

        # Use LLM to generate LlamaFactory config YAML
        # Coder will decide method based on hypothesis and available parameters
        config_files = self._generate_llamafactory_config_with_llm(
            base_model=base_model,
            task_info=task_info,
            queried_former_failed_knowledge=queried_former_failed_knowledge,
            prev_feedback=prev_task_feedback,
            workspace=workspace,
        )

        # Return generated config files directly - validation happens in evaluator
        return config_files

    def _generate_llamafactory_config_with_llm(
        self,
        base_model: str,
        task_info: str = "",
        queried_former_failed_knowledge: tuple = None,
        prev_feedback=None,
        workspace=None,
    ) -> dict[str, str]:
        """Generate LlamaFactory configuration YAML using LLM"""

        # Query LLaMA Factory parameters: shared params once + method-specific params
        available_methods = self.llama_factory_manager.methods
        shared_params = self.llama_factory_manager.format_shared_params()

        # Format method-specific parameters only (no duplication of shared params)
        methods_specific_params = {}
        for method in available_methods:
            methods_specific_params[method] = self.llama_factory_manager.format_method_specific_params(method)

        # Use environment-aware paths (Docker vs Conda)
        # Note: datasets_path in finetune_coder uses workspace path where processed
        # data.json and dataset_info.json are located (generated by FTDataEvaluator)

        # Generate prompts using templates with all required parameters
        system_prompt = T(".prompts:finetune_coder.system").r(
            scenario=self.scen.get_scenario_all_desc(),
            task_desc=task_info,
            queried_former_failed_knowledge=queried_former_failed_knowledge[0],
            available_methods=", ".join(available_methods),
            shared_params=shared_params,
            methods_specific_params=methods_specific_params,
        )

        # Read data_stats.json from workspace (injected by FTDataEvaluator)
        data_stats = workspace.file_dict.get("data_stats.json", "")

        user_prompt = T(".prompts:finetune_coder.user").r(
            latest_code=workspace.file_dict.get(FT_YAML_FILE_NAME, ""),
            latest_feedback=prev_feedback,
            base_model=base_model,
            models_path=FT_PATHS.models,
            datasets_path=FT_PATHS.workspace,  # Training config uses workspace path for processed data
            workspace_path=FT_PATHS.workspace,
            deepspeed_path=FT_PATHS.deepspeed,
            data_stats=data_stats,
            has_think_token=self.scen.model_info.get("has_think_token", False),
            force_think_token=FT_RD_SETTING.force_think_token,
        )

        # Call LLM to generate config (multi-turn)
        session = APIBackend().build_chat_session(session_system_prompt=system_prompt)

        # Turn 1: Generate main training config
        train_config_yaml = session.build_chat_completion(
            user_prompt=user_prompt,
            json_mode=False,
            code_block_language="yaml",
            code_block_fallback=False,
        )

        # Validate main config YAML syntax
        yaml.safe_load(train_config_yaml)
        logger.info("Extracted main YAML config successfully")

        # Turn 2: Generate test parameters (test_params.yaml)
        test_params_prompt = T(".prompts:finetune_coder.user_test_params").r(workspace_path=FT_PATHS.workspace)
        test_params_yaml = session.build_chat_completion(
            user_prompt=test_params_prompt,
            json_mode=False,
            code_block_language="yaml",
            code_block_fallback=False,
        )

        # Validate test params YAML syntax
        yaml.safe_load(test_params_yaml)
        logger.info("Extracted test params YAML successfully")

        return {FT_YAML_FILE_NAME: train_config_yaml, FT_TEST_PARAMS_FILE_NAME: test_params_yaml}


class LLMFinetuneCoSTEER(CoSTEER):
    """LLM Fine-tuning CoSTEER implementation"""

    def __init__(
        self,
        scen: Scenario,
        *args,
        **kwargs,
    ) -> None:
        settings = FTCoderCoSTEERSettings()
        eva = CoSTEERMultiEvaluator([FTDataEvaluator(scen=scen), FTCoderEvaluator(scen=scen)], scen=scen)
        es = LLMFinetuneEvolvingStrategy(scen=scen, settings=settings)

        super().__init__(
            *args,
            settings=settings,
            eva=eva,
            es=es,
            evolving_version=2,
            scen=scen,
            max_loop=FT_RD_SETTING.coder_max_loop if hasattr(FT_RD_SETTING, "coder_max_loop") else 5,
            stop_eval_chain_on_fail=True,  # finetune involve partial implementation.
            **kwargs,
        )