ai-nk-cce 0.1.0__py3-none-any.whl

model/nk_assistant.py

@@ -0,0 +1,207 @@
+import json
+from pathlib import Path
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from src.api.mpcdf_vllm import mpcdf_vllm_request
+from src.model.config import NKAssistantConfig, TrainConfig
+from src.model.parser import create_context, create_target, target_to_int
+
+
+def load_model_and_tokenizer(model_path: str):
+    """Load the model and tokenizer from the given path"""
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+    except Exception:
+        # Fix if tokenizer is not found
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    model = AutoModelForCausalLM.from_pretrained(model_path)
+
+    # Ensure pad_token is set
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Move model to GPU if available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+
+    return model, tokenizer, device
+
+
+class NKAssistant:
+    def __init__(
+        self, config: NKAssistantConfig, metadata: Optional[TrainConfig] = None
+    ):
+        if not config.use_mpcdf_vllm:
+            model, tokenizer, device = load_model_and_tokenizer(
+                config.model_path
+            )
+        else:
+            model = None
+            tokenizer = None
+            device = None
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = device
+        self.config = config
+        self.metadata = metadata
+        self.generation_params = config.generation_params
+
+    def _save_pretrained(self, save_directory: Path) -> None:
+        save_directory.mkdir(parents=True, exist_ok=True)
+        self.model.save_pretrained(save_directory)
+        self.tokenizer.save_pretrained(save_directory)
+        if self.metadata:
+            with open(save_directory / "metadata.json", "w") as f:
+                json.dump(self.metadata.model_dump(), f)
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: Union[str, Path], **kwargs
+    ):
+        # Use Hugging Face logic to locate the local or remote files
+        model_path = Path(pretrained_model_name_or_path)
+
+        if not model_path.exists():
+            # Download files from Hub into cache
+            from huggingface_hub import snapshot_download
+
+            model_path = Path(
+                snapshot_download(repo_id=pretrained_model_name_or_path)
+            )
+
+        metadata_path = model_path / "metadata.json"
+        metadata = None
+        if metadata_path.exists():
+            with open(metadata_path) as f:
+                metadata = TrainConfig(**json.load(f))
+
+        config = metadata.assistant_config
+        config.model_path = str(model_path)
+
+        return cls(config=config, metadata=metadata)
+
+    @staticmethod
+    def create_text_from_row(row, include_payoff=True, include_target=True):
+        sample_payoff = np.array(row["payoffs"])[row["sample_idx"]]
+
+        context = create_context(
+            n=row["n"],
+            k=row["k"],
+            power_scale=row["power_scale"],
+            sample_idxs=row["sample_idx"],
+            origin_idx=row["origin_idx"],
+            hamming_distance=row["hamming_distance"],
+            payoff=sample_payoff,
+            include_payoff=include_payoff,
+        )
+        if include_target:
+            target = create_target(target_idx=row["target_idx"])
+            return {"context": context, "target": target}
+        return {"context": context}
+
+    @staticmethod
+    def create_prompt_for_rl_from_row(row):
+        sample_payoff = np.array(row["payoffs"])[row["sample_idx"]]
+        context = create_context(
+            n=row["n"],
+            k=row["k"],
+            power_scale=row["power_scale"],
+            sample_idxs=row["sample_idx"],
+            origin_idx=row["origin_idx"],
+            hamming_distance=row["hamming_distance"],
+            payoff=sample_payoff,
+            include_payoff=True,
+        )
+        return {"prompt": context}
+
+    def suggest_target_from_context(
+        self, context: str, target_length: int
+    ) -> str:
+        # Tokenize input
+        input_ids = self.tokenizer.encode(context, return_tensors="pt").to(
+            self.device
+        )
+        n_input = input_ids.shape[1]
+
+        # Generate output
+        with torch.no_grad():
+            output = self.model.generate(
+                input_ids,
+                max_length=n_input + target_length,
+                pad_token_id=self.tokenizer.eos_token_id,
+                **self.generation_params,
+            )
+
+        # Decode only the new tokens (exclude the input tokens)
+        suggestion = self.tokenizer.decode(
+            output[0, n_input:], skip_special_tokens=True
+        )
+
+        return suggestion
+
+    def suggest_from_row(self, row):
+        context = self.create_text_from_row(row, include_target=False)[
+            "context"
+        ]
+        n = row["n"]
+        target_length = n * 2 - 1
+        if self.config.use_mpcdf_vllm:
+            binary_string = mpcdf_vllm_request(
+                prompt=context,
+                max_tokens=target_length,
+                **self.generation_params,
+            )
+        else:
+            binary_string = self.suggest_target_from_context(
+                context, target_length
+            )
+        return {"suggestion": binary_string}
+
+    def suggest(
+        self,
+        n: int,
+        k: int,
+        power_scale: float,
+        hamming_distance: int,
+        sample_idxs: List[int],
+        origin_idx: int,
+        payoffs: List[float],
+    ) -> int:
+        # Create context from input data
+        context = create_context(
+            n=n,
+            k=k,
+            power_scale=power_scale,
+            sample_idxs=sample_idxs,
+            origin_idx=origin_idx,
+            hamming_distance=hamming_distance,
+            payoff=payoffs,
+        )
+
+        # Set max suggestion length based on n
+        target_length = (
+            n * 2 - 1
+        )  # Length needed for binary string with commas
+
+        if self.config.use_mpcdf_vllm:
+            binary_string = mpcdf_vllm_request(
+                prompt=context,
+                max_tokens=target_length,
+                **self.generation_params,
+            )
+        else:
+            binary_string = self.suggest_target_from_context(
+                context, target_length
+            )
+
+        # Parse the suggestion to an integer
+        try:
+            suggestion_int = target_to_int(binary_string)
+            return suggestion_int
+        except Exception as e:
+            print(f"Error parsing suggestion '{binary_string}': {e}")
+            return -1

model/parser.py

@@ -0,0 +1,70 @@
+import logging
+from typing import List, Optional, Union
+
+from src.utils.binary_conversion import binary_str_to_int, int_to_binary_str
+
+logger = logging.getLogger()
+
+
+def create_context(
+    *,
+    n: int,
+    k: int,
+    power_scale: float,
+    sample_idxs: List[int],
+    origin_idx: int,
+    hamming_distance: int,
+    payoff: Optional[List[Union[float, int]]] = None,
+    include_payoff=True,
+    **kwargs,
+) -> str:
+    # Header
+    doc = ""
+    doc += f"n: {n}\n"
+    doc += f"k: {k}\n"
+    doc += f"p: {power_scale:.2f}\n"
+    doc += f"h: {hamming_distance}\n"
+    doc += "\n"
+    # Sample rows
+    if include_payoff:
+        assert (
+            payoff is not None
+        ), "Payoff values must be provided when include_payoff is True"
+        doc += "sample,payoff\n"
+        doc += "\n".join(
+            [
+                f"{int_to_binary_str(idx)},{p}"
+                for idx, p in zip(sample_idxs, payoff)
+            ]
+        )
+    else:
+        doc += "sample\n"
+        doc += "\n".join([int_to_binary_str(idx) for idx in sample_idxs])
+    doc += "\n"
+    doc += "\n"
+    # Constraints
+    doc += "user\n"
+    doc += int_to_binary_str(origin_idx) + "\n"
+    doc += "\n"
+    # Target
+    doc += "assistant\n"
+    return doc
+
+
+def create_target(*, target_idx: int) -> str:
+    """
+    Create a target string from a target index.
+    """
+    return int_to_binary_str(target_idx)
+
+
+def target_to_int(target):
+    """Convert target string to integer (deprecated, use binary_str_to_int)."""
+    int_target = binary_str_to_int(target)
+    logger.debug(f"target: {target}, int_target: {int_target}")
+    return int_target
+
+
+def eval_target(target, ranks):
+    target_idx = binary_str_to_int(target)
+    return ranks[target_idx]

model/run_slurm.py

@@ -0,0 +1,335 @@
+import argparse
+import os
+import re
+
+# import shutil
+import subprocess
+import sys
+from datetime import datetime
+
+# from collections import OrderedDict
+import yaml
+
+
+def generate_local_job_id():
+    """
+    Generates a local job ID based on timestamp.
+    """
+    return datetime.now().strftime("%Y_%m_%d__%H_%M_%S")
+
+
+def parse_unknown_args(unknown_args, argv):
+    # Process unknown arguments to create a dictionary
+    dynamic_args = {}
+    for arg in unknown_args:
+        if arg.startswith("--"):
+            key = arg.lstrip("-")
+            # Assuming the next item in the list is the value
+            if argv.index(arg) + 1 < len(argv):
+                value = argv[argv.index(arg) + 1]
+                dynamic_args[key] = value
+
+    return dynamic_args
+
+
+def define_output(args):
+    # Compute additional arguments
+    if args["output_dir"] is None:
+        args["output_dir"] = os.path.join(
+            "experiments",
+            args["group_name"],
+            args["job_name"],
+            args["job_id"],
+        )
+    if not os.path.exists(args["output_dir"]):
+        os.makedirs(args["output_dir"])
+    return args
+
+
+def define_compute_resources(args):
+    if args["n_gpu"] > 4:
+        assert args["n_gpu"] % 4 == 0
+        n_nodes = args["n_gpu"] // 4
+        n_gpu = 4
+    else:
+        n_nodes = 1
+        n_gpu = args["n_gpu"]
+
+    if n_gpu >= 4:
+        memory = 0
+    else:
+        memory = 125000 * n_gpu
+    partition = "gpu"
+    cpu = n_gpu * 18  # 18 cores per GPU
+
+    args = {
+        **args,
+        "n_nodes": n_nodes,
+        "n_gpu": n_gpu,
+        "n_cpu": cpu,
+        "partition": partition,
+        "memory": memory,
+    }
+    return args
+
+
+def deep_merge_configs(main_config, included_config):
+    """
+    Deep merges two configurations with the main config retaining its order.
+    """
+    for key, value in included_config.items():
+        if (
+            key in main_config
+            and isinstance(main_config[key], dict)
+            and isinstance(value, dict)
+        ):
+            deep_merge_configs(main_config[key], value)
+        else:
+            main_config[key] = value
+
+
+def find_include_value(data, target_key):
+    if isinstance(data, dict):
+        for key, value in data.items():
+            if key == target_key:
+                return value
+            elif isinstance(value, dict):
+                result = find_include_value(value, target_key)
+                if result is not None:
+                    return result
+    return None
+
+
+def _convert_to_float(value):
+    """Convert value to float, return None on error."""
+    try:
+        return float(value)
+    except ValueError:
+        return None
+
+
+def _convert_to_int(value):
+    """Convert value to int, return None on error."""
+    try:
+        return int(value)
+    except ValueError:
+        return None
+
+
+def _convert_to_list_int(value):
+    """Convert comma-separated string to list of ints, return None on error."""
+    try:
+        return list(map(int, value.split(",")))
+    except ValueError:
+        return None
+
+
+def _convert_replacement_value(placeholder_type, replacement_value):
+    """Convert replacement value to specified type."""
+    if replacement_value is None:
+        return replacement_value
+    if placeholder_type == "float":
+        return _convert_to_float(replacement_value)
+    if placeholder_type == "int":
+        return _convert_to_int(replacement_value)
+    if placeholder_type == "bool":
+        return replacement_value.lower() == "true"
+    if placeholder_type == "list_int":
+        return _convert_to_list_int(replacement_value)
+    return replacement_value
+
+
+def replace_placeholder(element, replacements=None):
+    if replacements is None:
+        replacements = {}
+
+    if isinstance(element, dict):
+        for key, value in element.items():
+            element[key] = replace_placeholder(value, replacements)
+    elif isinstance(element, list):
+        return [replace_placeholder(item, replacements) for item in element]
+    elif isinstance(element, str):
+
+        def replacement_function(match):
+            placeholder_type = match.group(1) if match.group(1) else "str"
+            placeholder_variable = match.group(2)
+            replacement_value = replacements.get(
+                placeholder_variable, match.group(0)
+            )
+            return _convert_replacement_value(
+                placeholder_type, replacement_value
+            )
+
+        pattern = r"<<(?:(\w+): )?(\w+)>>"
+        matches = list(re.finditer(pattern, element))
+
+        # If exactly one match and it spans the entire string, perform
+        # type conversion
+        if len(matches) == 1 and matches[0].span() == (0, len(element)):
+            return replacement_function(matches[0])
+
+        # For strings with multiple placeholders or additional text, replace
+        # without type conversion
+        def string_replacement_function(match):
+            return str(replacement_function(match))
+
+        result = re.sub(pattern, string_replacement_function, element)
+        return result
+
+    return element
+
+
+def load_and_merge_configs(
+    config_path,
+):
+    """
+    Loads configuration from the main file and merges included configurations
+    while preserving the order.
+    """
+    with open(config_path, "r") as file:
+        # Load the main configuration with FullLoader to preserve the order
+        main_config = yaml.load(file, Loader=yaml.FullLoader)
+
+    # Check if there are included configs and process them
+    includes = find_include_value(main_config, "__include")
+    if includes is not None:
+        for include_path in includes:
+            with open(include_path, "r") as inc_file:
+                included_config = yaml.load(inc_file, Loader=yaml.FullLoader)
+                deep_merge_configs(main_config, included_config)
+
+    return main_config
+
+
+def copy_config(config, args):
+    """
+    Copies the given config file to a job-specific directory after merging
+    included configurations. Preserves the order of parameters in the main
+    config file.
+    """
+    dest_filename = f"{args['job_id']}.yml"
+    dest_path = os.path.join(args["output_dir"], dest_filename)
+
+    with open(dest_path, "w") as file:
+        yaml.dump(
+            config, file, sort_keys=False
+        )  # Prevent sorting keys on dump
+    args["copied_config_file"] = dest_path
+    return args
+
+
+def generate_bash_script(args):
+    """
+    Reads in a bash template file and replaces placeholders with the given
+    config path. Writes the modified script to the job-specific directory.
+    """
+    output_path = os.path.join(args["output_dir"], f"{args['job_id']}.sh")
+
+    with open(args["template"], "r") as file:
+        script = file.read().format(**args)
+
+    with open(output_path, "w") as file:
+        file.write(script)
+    return output_path
+
+
+def submit_script(script_path):
+    """
+    Submits the given bash script to sbatch.
+    """
+    subprocess.run(["sbatch", script_path])
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Submit jobs with documentation of the YAML configuration."
+    )
+    parser.add_argument(
+        "--config_file",
+        type=str,
+        required=True,
+        help="Path to the YAML configuration file.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help=(
+            "Path to the output data. If both are None, path of config file "
+            "is used."
+        ),
+    )
+    parser.add_argument(
+        "--template",
+        type=str,
+        default="src/model/scripts/template.slurm",
+        help="Path to the bash script template.",
+    )
+    parser.add_argument(
+        "--dry",
+        action="store_true",
+        help="Only create files, do not submit the job.",
+    )
+    parser.add_argument(
+        "--n_gpu", type=int, default=1, help="Number of GPUs to use."
+    )
+    parser.add_argument(
+        "--time",
+        type=str,
+        default="00:10:00",
+        help="Expected runtime in HH:MM:SS format.",
+    )
+    parser.add_argument(
+        "--script",
+        type=str,
+        default="src/model/train.py",
+        help="Script to run.",
+    )
+    parser.add_argument(
+        "--job_name", type=str, default="v1", help="Name of the job."
+    )
+    parser.add_argument(
+        "--group_name",
+        type=str,
+        default="debug",
+        help="Name of the experiment group.",
+    )
+    parser.add_argument(
+        "--project_name",
+        type=str,
+        default="NK-Landscape",
+        help="Project to charge.",
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        default="/u/lumi/projects/llm-strategic-tuning/images/ai_nk_rl.sif",
+        help="Apptainer image to use",
+    )
+    argv = sys.argv[1:]
+
+    known_args, unknown_args = parser.parse_known_args(sys.argv[1:])
+    dynamic_args = parse_unknown_args(unknown_args, argv)
+    args_dict = vars(known_args)
+    args_dict.update(dynamic_args)
+
+    args_dict = define_compute_resources(args_dict)
+    args_dict["job_id"] = generate_local_job_id()
+    args_dict = args_dict.copy()
+
+    args_dict = define_output(args_dict)
+    config = load_and_merge_configs(args_dict["config_file"])
+    config = replace_placeholder(config, args_dict)
+    args_dict = copy_config(config, args_dict)
+
+    args_dict["config_file"] = args_dict["copied_config_file"]
+    generated_script = generate_bash_script(args_dict)
+
+    if not args_dict["dry"]:
+        submit_script(generated_script)
+    else:
+        print(f"Generated script at {generated_script} without submission.")
+
+
+if __name__ == "__main__":
+    main()