From 74e0d811842603d4e8c8c9aebd03d8023d17eac5 Mon Sep 17 00:00:00 2001 From: mcclain Date: Sun, 26 Oct 2025 08:54:50 +0000 Subject: [PATCH 01/11] slow --- docker-compose.yaml | 2 ++ src/config.py | 3 +++ src/rewards/reward_config.py | 33 +++++++++++++++++++++++++++ src/rewards/scorer.py | 44 ++++++++++++++++++++++++++++++++++++ src/runners/grpo.py | 32 +++++++++++++++++++------- 5 files changed, 106 insertions(+), 8 deletions(-) create mode 100644 src/rewards/reward_config.py create mode 100644 src/rewards/scorer.py diff --git a/docker-compose.yaml b/docker-compose.yaml index 32c0544..00ca6c1 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -36,6 +36,7 @@ services: - /mnt/s3/phd-research-storage-1758274488:/s3:rw command: > bash -lc " + mkdir -p /tmp/wandb && \ uv sync --frozen && uv run python -m src.runners.grpo " environment: @@ -44,6 +45,7 @@ services: - WANDB_PROJECT=plasmidrl-trl-grpo - WANDB_TAGS=["plasmid","rl","trl","grpo"] - WANDB_NOTES=TRL GRPO training on plasmid design + - WANDB_DIR=/tmp/wandb # Single simple flag to optionally resume - RESUME=${RESUME:-false} diff --git a/src/config.py b/src/config.py index 71de066..872d0b2 100644 --- a/src/config.py +++ b/src/config.py @@ -20,6 +20,9 @@ class Config(BaseSettings): validation_alias=AliasChoices("hf_token", "HF_TOKEN", "HUGGINGFACE_TOKEN"), ) + train_dataset: str = "data/train.parquet" + val_dataset: str = "data/test.parquet" + #this is the GFP cassette default_query: str = "tttacggctagctcagtcctaggtatagtgctagcTACTagagaaagaggagaaatactaAATGatgcgtaaaggagaagaacttttcactggagttgtcccaattcttgttgaattagatggtgatgttaatgggcacaaattttctgtcagtggagagggtgaaggtgatgcaacatacggaaaacttacccttaaatttatttgcactactggaaaactacctgttccatggccaacacttgtcactactttcggttatggtgttcaatgctttgcgagatacccagatcatatgaaacagcatgactttttcaagagtgccatgcccgaaggttatgtacaggaaagaactatatttttcaaagatgacgggaactacaagacacgtgctgaagtcaagtttgaaggtgatacccttgttaatagaatcgagttaaaaggtattgattttaaagaagatggaaacattcttggacacaaattggaatacaactataactcacacaatgtatacatcatggcagacaaacaaaagaatggaatcaaagttaacttcaaaattagacacaacattgaagatggaagcgttcaactagcagaccattatcaacaaaatactccaattggcgatggccctgtccttttaccagacaaccattacctgtccacacaatctgccctttcgaaagatcccaacgaaaagagagatcacatggtccttcttgagtttgtaacagctgttgtttgtcggtgaacgctctctactagagtcacactggctcaccttcgggtgggcctttctgcgtttata".upper() diff --git a/src/rewards/reward_config.py b/src/rewards/reward_config.py new file mode 100644 index 0000000..a4144c4 --- /dev/null +++ b/src/rewards/reward_config.py @@ -0,0 +1,33 @@ +from pydantic import BaseModel + +class RewardConfig(BaseModel): + + punish_mode: bool = True # penaize violations of the reward config as opposed to just not rewarding them + length_penalty: bool = True # penalize sequences that are too long or too short + location_aware: bool = True # reward sequences that are located in the correct location (e.g. prompoter then cds then terminatr) + + ori_min: int = 1 + ori_max: int = 2 + allowed_oris: Optional[List[str]] = None + ori_weight: float = 1.0 + + promoter_min: int = 1 + promoter_max: int = 2 + allowed_promoters: Optional[List[str]] = None + promoter_weight: float = 1.0 + + terminator_min: int = 1 + terminator_max: int = 2 + allowed_terminators: Optional[List[str]] = None + terminator_weight: float = 1.0 + + marker_min: int = 1 + marker_max: int = 2 + allowed_markers: Optional[List[str]] = None + marker_weight: float = 1.0 + + cds_min: int = 1 + cds_max: int = 2 + allowed_cds: Optional[List[str]] = None + cds_weight: float = 1.0 + diff --git a/src/rewards/scorer.py b/src/rewards/scorer.py new file mode 100644 index 0000000..6f60ef9 --- /dev/null +++ b/src/rewards/scorer.py @@ -0,0 +1,44 @@ +from rewards.reward_config import RewardConfig +import plasmidkit as pk +import pandas as pd + + +class Scorer: + def __init__(self, reward_config: RewardConfig): + self.reward_config = reward_config + self.score_functions = [self.score_ori, self.score_promoter, self.score_terminator, self.score_marker, self.score_cds] + self.weights = [self.reward_config.ori_weight, self.reward_config.promoter_weight, self.reward_config.terminator_weight, self.reward_config.marker_weight, self.reward_config.cds_weight] + + def annotate(self, sequence: str) -> Any: + return pk.annotate(sequence, is_sequence=True) + + def runner(self, sequence: str, annotations: Any) -> float: + """ + Takes a list of runnable functions with the same signature and runs them in parallel. + + Args: + sequence: The sequence to score + annotations: The annotations to score + + Returns: + The score + """ + with ThreadPoolExecutor(max_workers=len(self.score_functions)) as executor: + futures = [executor.submit(func, sequence, annotations) for func in self.score_functions] + return [future.result() for future in futures] + pass + + def score_ori(self, annotations: pd.DataFrame) -> float: + return 0.0 + + def score_promoter(self, annotations: pd.DataFrame) -> float: + return 0.0 + + def score_terminator(self, annotations: pd.DataFrame) -> float: + return 0.0 + + def score_marker(self, annotations: pd.DataFrame) -> float: + return 0.0 + + def score_cds(self, annotations: pd.DataFrame) -> float: + return 0.0 \ No newline at end of file diff --git a/src/runners/grpo.py b/src/runners/grpo.py index 9057eeb..197dd31 100644 --- a/src/runners/grpo.py +++ b/src/runners/grpo.py @@ -36,16 +36,32 @@ def select_prompt_and_extras(ds): eval_ds = None use_eval = False -# ---- tokenizer ---- + tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True, trust_remote_code=True) -if tok.pad_token is None: - tok.pad_token = tok.eos_token tok.padding_side = "left" +# Remap specials to the correct strings (these already exist in your vocab) +tok.eos_token = "" +tok.bos_token = "" +tok.pad_token = "[PAD]" + +# Re-assert +assert tok.eos_token_id == 30001, tok.eos_token_id +assert tok.bos_token_id == 30000, tok.bos_token_id +assert tok.pad_token_id == 3, tok.pad_token_id + +# Pass IDs explicitly to the model so nothing “helpfully” changes at runtime +model_init_kwargs = { + "trust_remote_code": True, + "eos_token_id": tok.eos_token_id, + "bos_token_id": tok.bos_token_id, + "pad_token_id": tok.pad_token_id, +} # ---- GRPO config (keys verified against TRL docs) ---- args = GRPOConfig( + model_init_kwargs=model_init_kwargs, # transformers-style output_dir=save_path, num_train_epochs=20, @@ -56,7 +72,7 @@ def select_prompt_and_extras(ds): gradient_accumulation_steps=1, max_steps=-1, save_strategy="steps", - save_steps=10, + save_steps=100, logging_strategy="steps", logging_steps=1, bf16=torch.cuda.is_available(), @@ -65,17 +81,17 @@ def select_prompt_and_extras(ds): seed=SEED, do_eval=use_eval, eval_strategy="steps" if use_eval else "no", - eval_steps=10 if use_eval else None, + eval_steps=100 if use_eval else None, # model/ref-model handling - model_init_kwargs={"trust_remote_code": True}, # used if model is passed as string disable_dropout=True, # stabilizes ref-policy logprobs # data & generation remove_unused_columns=False, # keep your extras for reward_fn max_prompt_length=1024, num_generations=8, - max_completion_length=256, + max_completion_length=256, + #min_completion_length=16, #not supported currently temperature=0.80, top_p=0.90, @@ -84,7 +100,7 @@ def select_prompt_and_extras(ds): epsilon=0.2, # PPO-style clip (replaces cliprange) loss_type="bnpo", # token-level normalization; avoids length bias scale_rewards=True, - mask_truncated_completions=True, + mask_truncated_completions=False, # vLLM (colocated serverless flag is not a key here) use_vllm=True, From 9c06457329cc3ab3d46a8e43faf132fd45df63e6 Mon Sep 17 00:00:00 2001 From: mcclain Date: Sun, 26 Oct 2025 09:26:53 +0000 Subject: [PATCH 02/11] hm --- src/rewards/__init__.py | 18 +++++++----------- src/rewards/scorer.py | 32 +++++++++++++++++++++++--------- 2 files changed, 30 insertions(+), 20 deletions(-) diff --git a/src/rewards/__init__.py b/src/rewards/__init__.py index d65af25..f82eada 100644 --- a/src/rewards/__init__.py +++ b/src/rewards/__init__.py @@ -1,12 +1,8 @@ -from .rewards import score_completions -# The reward functions are expected to accept: +from .scorer import Scorer +from .reward_config import RewardConfig -# completions: list of model outputs for the batch -# other keyword arguments corresponding to columns in your dataset (e.g. solution, ground_truth, image, etc.) -# Optionally, **kwargs to absorb extra parameters (like trainer_state) -# this always has the same signature so I just define it here once. -def Score( - completions: list, - **kwargs, -) -> list[float]: - return score_completions(completions) \ No newline at end of file +config = RewardConfig( + +) + +score = Scorer(RewardConfig()).score \ No newline at end of file diff --git a/src/rewards/scorer.py b/src/rewards/scorer.py index 6f60ef9..2d617f4 100644 --- a/src/rewards/scorer.py +++ b/src/rewards/scorer.py @@ -1,6 +1,8 @@ from rewards.reward_config import RewardConfig import plasmidkit as pk import pandas as pd +from typing import Any +from concurrent.futures import ThreadPoolExecutor class Scorer: @@ -8,11 +10,17 @@ def __init__(self, reward_config: RewardConfig): self.reward_config = reward_config self.score_functions = [self.score_ori, self.score_promoter, self.score_terminator, self.score_marker, self.score_cds] self.weights = [self.reward_config.ori_weight, self.reward_config.promoter_weight, self.reward_config.terminator_weight, self.reward_config.marker_weight, self.reward_config.cds_weight] + total_weight = sum(self.weights) + if total_weight: + self.weights = [w / total_weight for w in self.weights] + else: + uniform_weight = 1.0 / len(self.score_functions) + self.weights = [uniform_weight for _ in self.score_functions] def annotate(self, sequence: str) -> Any: return pk.annotate(sequence, is_sequence=True) - def runner(self, sequence: str, annotations: Any) -> float: + def runner(self, sequence: str, annotations: Any) -> list[float]: """ Takes a list of runnable functions with the same signature and runs them in parallel. @@ -21,24 +29,30 @@ def runner(self, sequence: str, annotations: Any) -> float: annotations: The annotations to score Returns: - The score + A list of scores in the same order as `self.score_functions`. """ with ThreadPoolExecutor(max_workers=len(self.score_functions)) as executor: futures = [executor.submit(func, sequence, annotations) for func in self.score_functions] return [future.result() for future in futures] - pass - def score_ori(self, annotations: pd.DataFrame) -> float: + def score_ori(self, seq: str, annotations: pd.DataFrame) -> float: return 0.0 - def score_promoter(self, annotations: pd.DataFrame) -> float: + def score_promoter(self, seq: str, annotations: pd.DataFrame) -> float: return 0.0 - def score_terminator(self, annotations: pd.DataFrame) -> float: + def score_terminator(self, seq: str, annotations: pd.DataFrame) -> float: return 0.0 - def score_marker(self, annotations: pd.DataFrame) -> float: + def score_marker(self, seq: str, annotations: pd.DataFrame) -> float: return 0.0 - def score_cds(self, annotations: pd.DataFrame) -> float: - return 0.0 \ No newline at end of file + def score_cds(self, seq: str, annotations: pd.DataFrame) -> float: + return 0.0 + + def score_length(self, seq: str, annotations: pd.DataFrame) -> float: + return 0.0 + + def score(self, seq: str, annotations: pd.DataFrame) -> float: + results = self.runner(seq, annotations) + return sum(w * r for w, r in zip(self.weights, results)) \ No newline at end of file From aba210bbd3b9bdde5018efa1f0077acc1427c726 Mon Sep 17 00:00:00 2001 From: McClain Thiel Date: Fri, 31 Oct 2025 12:09:16 +0000 Subject: [PATCH 03/11] working better i think --- .DS_Store | Bin 0 -> 8196 bytes notebooks/ray_test.ipynb | 93 ++++++++++ pyproject.toml | 1 + src/eval/infer.py | 46 +++++ src/rewards/reward_config.py | 30 +++- src/rewards/rewards.py | 209 +++------------------- src/rewards/scorer.py | 335 ++++++++++++++++++++++++++++++++--- tests/data/RF0G-IodoY.fasta | 79 +++++++++ tests/data/pSC101.fasta | 2 + tests/data/pUC19.fasta | 2 + tests/test_rewards.py | 98 ++++++++++ uv.lock | 40 ++++- 12 files changed, 721 insertions(+), 214 deletions(-) create mode 100644 .DS_Store create mode 100644 notebooks/ray_test.ipynb create mode 100644 src/eval/infer.py create mode 100644 tests/data/RF0G-IodoY.fasta create mode 100644 tests/data/pSC101.fasta create mode 100644 tests/data/pUC19.fasta create mode 100644 tests/test_rewards.py diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..1f447099abcc9847609503430c5335e257df3fec GIT binary patch literal 8196 zcmeHMzl#$=6n>N2J0n84M!9koD+}=|3sKgvSBRn~*Vvg$az-(6LvojT!NyAe06{^q z*3wQJxms-O6n_(HZ@w?zB-xz}5t;f{dzom7h^pv} zi-$0jG>&sGw2GfO4m7BzZfAL|mG%0COr1~%)B$xs9Z(0eRzVO%-pW1p9J z%pJONGXC&kJhSm96l2-3er(gpBo1{{2h@S01AKOm=p=RM7H!J&yLovcm9buyrp<2F z#9a0HhsN1QtJ~+pczf4)|K{Y`fO-1@r8E!T5pB>N8cF8tXzgwomE6Z3Qf2$Hw@>+*GlfPc*y&cc^AN&!Gqh_gF(n$do~@H{NElw&F{PY57v+x zX61xB64o&f|9KY%(+-@fg;A;DQ{H^_@rwT(jO(R>@{Ewb!Tblop27bOY63I(F$~)YwR>iuuabb64$u(0G=NKes0_8ykuSvzsq5>LA6Fl^y_#A_ zYMPxD=&`o}GzUEcY6Tvd)9<0fM+yYB{{$H8>{C}@BTCWbM1N+zk6Qyfu9RmHewIG$x+A(?qog3%n4s8hr hPT~2; \u001b[39m\u001b[32m1\u001b[39m \u001b[43mray\u001b[49m\u001b[43m.\u001b[49m\u001b[43minit\u001b[49m\u001b[43m(\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mray://localhost:10001\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[32m 3\u001b[39m \u001b[38;5;129m@ray\u001b[39m.remote\n\u001b[32m 4\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mdouble\u001b[39m(x): \n\u001b[32m 5\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m x * \u001b[32m2\u001b[39m\n", + "\u001b[36mFile \u001b[39m\u001b[32m~/Projects/PhD/PlasmidRL/.venv/lib/python3.11/site-packages/ray/_private/client_mode_hook.py:104\u001b[39m, in \u001b[36mclient_mode_hook..wrapper\u001b[39m\u001b[34m(*args, **kwargs)\u001b[39m\n\u001b[32m 102\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m func.\u001b[34m__name__\u001b[39m != \u001b[33m\"\u001b[39m\u001b[33minit\u001b[39m\u001b[33m\"\u001b[39m \u001b[38;5;129;01mor\u001b[39;00m is_client_mode_enabled_by_default:\n\u001b[32m 103\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mgetattr\u001b[39m(ray, func.\u001b[34m__name__\u001b[39m)(*args, **kwargs)\n\u001b[32m--> \u001b[39m\u001b[32m104\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[36mFile \u001b[39m\u001b[32m~/Projects/PhD/PlasmidRL/.venv/lib/python3.11/site-packages/ray/_private/worker.py:1655\u001b[39m, in \u001b[36minit\u001b[39m\u001b[34m(address, num_cpus, num_gpus, resources, labels, object_store_memory, local_mode, ignore_reinit_error, include_dashboard, dashboard_host, dashboard_port, job_config, configure_logging, logging_level, logging_format, logging_config, log_to_driver, namespace, runtime_env, enable_resource_isolation, system_reserved_cpu, system_reserved_memory, **kwargs)\u001b[39m\n\u001b[32m 1653\u001b[39m passed_kwargs.update(kwargs)\n\u001b[32m 1654\u001b[39m builder._init_args(**passed_kwargs)\n\u001b[32m-> \u001b[39m\u001b[32m1655\u001b[39m ctx = \u001b[43mbuilder\u001b[49m\u001b[43m.\u001b[49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 1656\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mray\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01m_common\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01musage\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m usage_lib\n\u001b[32m 1658\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m passed_kwargs.get(\u001b[33m\"\u001b[39m\u001b[33mallow_multiple\u001b[39m\u001b[33m\"\u001b[39m) \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n", + "\u001b[36mFile \u001b[39m\u001b[32m~/Projects/PhD/PlasmidRL/.venv/lib/python3.11/site-packages/ray/client_builder.py:173\u001b[39m, in \u001b[36mClientBuilder.connect\u001b[39m\u001b[34m(self)\u001b[39m\n\u001b[32m 170\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m._allow_multiple_connections:\n\u001b[32m 171\u001b[39m old_ray_cxt = ray.util.client.ray.set_context(\u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[32m--> \u001b[39m\u001b[32m173\u001b[39m client_info_dict = \u001b[43mray\u001b[49m\u001b[43m.\u001b[49m\u001b[43mutil\u001b[49m\u001b[43m.\u001b[49m\u001b[43mclient_connect\u001b[49m\u001b[43m.\u001b[49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m 174\u001b[39m \u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43maddress\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 175\u001b[39m \u001b[43m \u001b[49m\u001b[43mjob_config\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_job_config\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 176\u001b[39m \u001b[43m \u001b[49m\u001b[43m_credentials\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_credentials\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 177\u001b[39m \u001b[43m \u001b[49m\u001b[43mray_init_kwargs\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_remote_init_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 178\u001b[39m \u001b[43m \u001b[49m\u001b[43mmetadata\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_metadata\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 179\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 181\u001b[39m dashboard_url = ray.util.client.ray._get_dashboard_url()\n\u001b[32m 183\u001b[39m cxt = ClientContext(\n\u001b[32m 184\u001b[39m dashboard_url=dashboard_url,\n\u001b[32m 185\u001b[39m python_version=client_info_dict[\u001b[33m\"\u001b[39m\u001b[33mpython_version\u001b[39m\u001b[33m\"\u001b[39m],\n\u001b[32m (...)\u001b[39m\u001b[32m 189\u001b[39m _context_to_restore=ray.util.client.ray.get_context(),\n\u001b[32m 190\u001b[39m )\n", + "\u001b[36mFile \u001b[39m\u001b[32m~/Projects/PhD/PlasmidRL/.venv/lib/python3.11/site-packages/ray/util/client_connect.py:55\u001b[39m, in \u001b[36mconnect\u001b[39m\u001b[34m(conn_str, secure, metadata, connection_retries, job_config, namespace, ignore_version, _credentials, ray_init_kwargs)\u001b[39m\n\u001b[32m 50\u001b[39m _explicitly_enable_client_mode()\n\u001b[32m 52\u001b[39m \u001b[38;5;66;03m# TODO(barakmich): https://github.com/ray-project/ray/issues/13274\u001b[39;00m\n\u001b[32m 53\u001b[39m \u001b[38;5;66;03m# for supporting things like cert_path, ca_path, etc and creating\u001b[39;00m\n\u001b[32m 54\u001b[39m \u001b[38;5;66;03m# the correct metadata\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m55\u001b[39m conn = \u001b[43mray\u001b[49m\u001b[43m.\u001b[49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m 56\u001b[39m \u001b[43m \u001b[49m\u001b[43mconn_str\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 57\u001b[39m \u001b[43m \u001b[49m\u001b[43mjob_config\u001b[49m\u001b[43m=\u001b[49m\u001b[43mjob_config\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 58\u001b[39m \u001b[43m \u001b[49m\u001b[43msecure\u001b[49m\u001b[43m=\u001b[49m\u001b[43msecure\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 59\u001b[39m \u001b[43m \u001b[49m\u001b[43mmetadata\u001b[49m\u001b[43m=\u001b[49m\u001b[43mmetadata\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 60\u001b[39m \u001b[43m \u001b[49m\u001b[43mconnection_retries\u001b[49m\u001b[43m=\u001b[49m\u001b[43mconnection_retries\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 61\u001b[39m \u001b[43m \u001b[49m\u001b[43mnamespace\u001b[49m\u001b[43m=\u001b[49m\u001b[43mnamespace\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 62\u001b[39m \u001b[43m \u001b[49m\u001b[43mignore_version\u001b[49m\u001b[43m=\u001b[49m\u001b[43mignore_version\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 63\u001b[39m \u001b[43m \u001b[49m\u001b[43m_credentials\u001b[49m\u001b[43m=\u001b[49m\u001b[43m_credentials\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 64\u001b[39m \u001b[43m \u001b[49m\u001b[43mray_init_kwargs\u001b[49m\u001b[43m=\u001b[49m\u001b[43mray_init_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 65\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 66\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m conn\n", + "\u001b[36mFile \u001b[39m\u001b[32m~/Projects/PhD/PlasmidRL/.venv/lib/python3.11/site-packages/ray/util/client/__init__.py:233\u001b[39m, in \u001b[36mRayAPIStub.connect\u001b[39m\u001b[34m(self, *args, **kw_args)\u001b[39m\n\u001b[32m 231\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mconnect\u001b[39m(\u001b[38;5;28mself\u001b[39m, *args, **kw_args):\n\u001b[32m 232\u001b[39m \u001b[38;5;28mself\u001b[39m.get_context()._inside_client_test = \u001b[38;5;28mself\u001b[39m._inside_client_test\n\u001b[32m--> \u001b[39m\u001b[32m233\u001b[39m conn = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mget_context\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m.\u001b[49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkw_args\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 234\u001b[39m \u001b[38;5;28;01mglobal\u001b[39;00m _lock, _all_contexts\n\u001b[32m 235\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m _lock:\n", + "\u001b[36mFile \u001b[39m\u001b[32m~/Projects/PhD/PlasmidRL/.venv/lib/python3.11/site-packages/ray/util/client/__init__.py:99\u001b[39m, in \u001b[36m_ClientContext.connect\u001b[39m\u001b[34m(self, conn_str, job_config, secure, metadata, connection_retries, namespace, ignore_version, _credentials, ray_init_kwargs)\u001b[39m\n\u001b[32m 97\u001b[39m \u001b[38;5;28mself\u001b[39m.client_worker._server_init(job_config, ray_init_kwargs)\n\u001b[32m 98\u001b[39m conn_info = \u001b[38;5;28mself\u001b[39m.client_worker.connection_info()\n\u001b[32m---> \u001b[39m\u001b[32m99\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_check_versions\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconn_info\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mignore_version\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 100\u001b[39m \u001b[38;5;28mself\u001b[39m._register_serializers()\n\u001b[32m 101\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m conn_info\n", + "\u001b[36mFile \u001b[39m\u001b[32m~/Projects/PhD/PlasmidRL/.venv/lib/python3.11/site-packages/ray/util/client/__init__.py:121\u001b[39m, in \u001b[36m_ClientContext._check_versions\u001b[39m\u001b[34m(self, conn_info, ignore_version)\u001b[39m\n\u001b[32m 118\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m_check_versions\u001b[39m(\u001b[38;5;28mself\u001b[39m, conn_info: Dict[\u001b[38;5;28mstr\u001b[39m, Any], ignore_version: \u001b[38;5;28mbool\u001b[39m) -> \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m 119\u001b[39m \u001b[38;5;66;03m# conn_info has \"python_version\" and \"ray_version\" so it can be used to compare.\u001b[39;00m\n\u001b[32m 120\u001b[39m ignore_version = ignore_version \u001b[38;5;129;01mor\u001b[39;00m (\u001b[33m\"\u001b[39m\u001b[33mRAY_IGNORE_VERSION_MISMATCH\u001b[39m\u001b[33m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m os.environ)\n\u001b[32m--> \u001b[39m\u001b[32m121\u001b[39m \u001b[43mcheck_version_info\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m 122\u001b[39m \u001b[43m \u001b[49m\u001b[43mconn_info\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 123\u001b[39m \u001b[43m \u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mRay Client\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m 124\u001b[39m \u001b[43m \u001b[49m\u001b[43mraise_on_mismatch\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mignore_version\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m 125\u001b[39m \u001b[43m \u001b[49m\u001b[43mpython_version_match_level\u001b[49m\u001b[43m=\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mminor\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m 126\u001b[39m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[36mFile \u001b[39m\u001b[32m~/Projects/PhD/PlasmidRL/.venv/lib/python3.11/site-packages/ray/_private/utils.py:1286\u001b[39m, in \u001b[36mcheck_version_info\u001b[39m\u001b[34m(cluster_metadata, this_process_address, raise_on_mismatch, python_version_match_level)\u001b[39m\n\u001b[32m 1284\u001b[39m error_message = \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mVersion mismatch: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mmismatch_msg\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m\n\u001b[32m 1285\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m raise_on_mismatch:\n\u001b[32m-> \u001b[39m\u001b[32m1286\u001b[39m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(error_message)\n\u001b[32m 1287\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m 1288\u001b[39m logger.warning(error_message)\n", + "\u001b[31mRuntimeError\u001b[39m: Version mismatch: The cluster was started with:\n Ray: 2.49.2\n Python: 3.9.18\nThis process on Ray Client was started with:\n Ray: 2.49.2\n Python: 3.11.9\n" + ] + } + ], + "source": [ + "ray.init(\"ray://localhost:10001\")\n", + "\n", + "@ray.remote\n", + "def double(x): \n", + " return x * 2\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "69fe2208", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index 9586827..5110d50 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ dependencies = [ "psutil>=7.0.0", "pydantic>=2.0.0", "pydantic-settings>=2.0.0", + "pytest>=8.4.2", "pyzmq>=27.0.2", "tensordict>=0.9.1", "torchrl[llm]>=0.9.2", diff --git a/src/eval/infer.py b/src/eval/infer.py new file mode 100644 index 0000000..37db0eb --- /dev/null +++ b/src/eval/infer.py @@ -0,0 +1,46 @@ +import boto3 +import ray +from vllm import LLM, SamplingParams +import pandas as pd +import datetime + +from src.config import get_config + + +config = get_config() + +def main(): + + + ds = ray.data.read_parquet("s3://anonymous@air-example-data/prompts.txt") + + gfp = ray.data.read_text("s3://phd-research-storage-1758274488/prompts/GFP_cassette.fasta") + gfp_cassette = "".join([x['text'].upper() for x in ds.iter_rows()][1:]) + + prompts = [ + gfp_cassette, + "ATG" + ] * 50 + + sampling_params = SamplingParams(temperature=0.7, top_p=0.95) + + llm = LLM(model=config.model, tokenizer=config.tokenizer) + + results = llm.generate(prompts, sampling_params) + + records = [] + for output in results: + records.append({ + "prompt": output.prompt, + "response": output.outputs[0].text, + "full": output.prompt + output.outputs[0].text, + "length": len(output.prompt + output.outputs[0].text), + }) + + df = pd.DataFrame(records) + id = f"{config.model}-{datetime.now().strftime('%Y%m%d%H%M%S')}" + df.to_parquet(config.storage_bucket + "infer.parquet") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/rewards/reward_config.py b/src/rewards/reward_config.py index a4144c4..b91aecc 100644 --- a/src/rewards/reward_config.py +++ b/src/rewards/reward_config.py @@ -1,25 +1,30 @@ from pydantic import BaseModel +from typing import Optional, List class RewardConfig(BaseModel): punish_mode: bool = True # penaize violations of the reward config as opposed to just not rewarding them length_penalty: bool = True # penalize sequences that are too long or too short + min_length: Optional[int] = None + max_length: Optional[int] = None location_aware: bool = True # reward sequences that are located in the correct location (e.g. prompoter then cds then terminatr) + # Penalty factor applied when min/max constraints are violated (outside of range) + violation_penalty_factor: float = 0.5 ori_min: int = 1 - ori_max: int = 2 + ori_max: int = 1 allowed_oris: Optional[List[str]] = None ori_weight: float = 1.0 promoter_min: int = 1 - promoter_max: int = 2 + promoter_max: int = 1 allowed_promoters: Optional[List[str]] = None promoter_weight: float = 1.0 - terminator_min: int = 1 + terminator_min: int = 0 terminator_max: int = 2 allowed_terminators: Optional[List[str]] = None - terminator_weight: float = 1.0 + terminator_weight: float = 0.5 marker_min: int = 1 marker_max: int = 2 @@ -27,7 +32,22 @@ class RewardConfig(BaseModel): marker_weight: float = 1.0 cds_min: int = 1 - cds_max: int = 2 + cds_max: int = 5 allowed_cds: Optional[List[str]] = None cds_weight: float = 1.0 + # Location-aware cassette scoring constants (simplified) + cassette_max_cassettes: int = 2 + cassette_order_points: float = 5.0 + cassette_proximity_points: float = 5.0 + cassette_max_points_per: float = 20.0 + # Single proximity threshold (bp) for awarding proximity points + proximity_threshold_bp: int = 300 + # Final CDS location-aware bonus scale (added on top of count score, then clamped) + location_bonus_scale: float = 0.5 + # Overlap merge threshold (fraction of the smaller interval that must overlap) + overlap_merge_threshold: float = 0.8 + + def log_to_wandb(self): + raise NotImplementedError("Not implemented") + diff --git a/src/rewards/rewards.py b/src/rewards/rewards.py index 4dcbb35..a60f2b8 100644 --- a/src/rewards/rewards.py +++ b/src/rewards/rewards.py @@ -1,21 +1,26 @@ from concurrent.futures import ThreadPoolExecutor from functools import partial -from typing import Any -import plasmidkit as pk +from typing import Any, List import logging import time -from typing import List, Tuple, Any, Iterable + +import plasmidkit as pk + +from src.rewards.scorer import Scorer +from src.rewards.reward_config import RewardConfig logger = logging.getLogger("reward_logger") -# Toggle detailed timing logs from Config if available +# Toggle detailed timing logs try: - from src.config import Config # local import to avoid heavy deps at import time + from src.config import Config _REWARD_LOG_TIMINGS = bool(Config().reward_log_timings) except Exception: _REWARD_LOG_TIMINGS = False +_SCORER = Scorer(RewardConfig()) + def annotate_completions(completions: list[str]) -> list[Any]: """Annotate a flat list of completions using threads; strips spaces from sequences.""" @@ -31,204 +36,42 @@ def annotate_completions(completions: list[str]) -> list[Any]: return annotations - -def score_sequence( - sequence: str, - annotations: Any, - target_keywords: Iterable[str] | None = None, -) -> float: - """ - Heuristic backbone score for GRPO. - - Rewards: - • Single ORI (+20; multiple get +10 then −5 per extra, capped). - • Up to two highest-scoring cassettes with partial credit: - - promoter→CDS: order (+5) + proximity (≤100bp:+5, ≤300:+3, ≤500:+2, ≤1000:+1). - - CDS→terminator: order (+5) + proximity (same as above). - - Out-of-order legs get proximity-only partial (+≤3). - - +2 if CDS is a marker and promoter is within 300 bp. - • Standalone promoters (+1 each up to +5). - • Standalone terminators (+1 each up to +5). - • Payload (GOI) CDS anywhere: +8, plus +4 if any promoter within 500 bp. - - Penalty: - • Piecewise length penalty (1–10 kb, max −15) with **gentler slope at 3–5 kb**: - 1–3 kb: up to −3; 3–5 kb: to −5; 5–7.5 kb: to −9; 7.5–10 kb: to −15. - Ensures removing useful features can't increase the score. - - Returns [0, 100]. - """ - # ---- Early return for empty sequence ---- - if not sequence or len(sequence.strip()) == 0: +def score_sequence(sequence: str) -> float: + """Config-driven score in [0, 100] using Scorer (auto-annotates).""" + if not sequence or not str(sequence).strip(): return 0.0 - - # ---- collect features (case-insensitive 'type') ---- - def T(x): return (x.type or "").lower() - feats = list(annotations) - oris = [x for x in feats if T(x) in ("rep_origin", "ori", "origin_of_replication")] - promoters = [x for x in feats if T(x) == "promoter"] - cdss = [x for x in feats if T(x) == "cds"] - terminators = [x for x in feats if T(x) == "terminator"] - - # ---- small helpers ---- - def strand(x): return getattr(x, "strand", "+") - def start(x): return int(getattr(x, "start", 0)) - def end(x): return int(getattr(x, "end", 0)) - def role(x): return (getattr(x, "evidence", {}) or {}).get("role", "").lower() - - def in_order_same_strand(a, b) -> bool: - if strand(a) != strand(b): return False - if strand(a) == "+": return start(a) <= start(b) - else: return end(a) >= end(b) - - def distance(a, b) -> int: - if end(a) < start(b): return start(b) - end(a) - if end(b) < start(a): return start(a) - end(b) - return 0 # overlapping/adjacent - - def prox_points(d: int, max_pts: int = 5) -> int: - if d <= 100: return max_pts - if d <= 300: return max_pts - 2 # 3 - if d <= 500: return max_pts - 3 # 2 - if d <= 1000: return 1 - return 0 - - # ---- ORI scoring ---- - score = 0.0 - if len(oris) == 0: - pass - elif len(oris) == 1: - score += 20.0 - else: - score += 10.0 - score -= min(15.0, 5.0 * (len(oris) - 1)) - - # ---- Cassette scoring (incl. out-of-order partials) ---- - def best_cassettes() -> List[Tuple[Any, Any, Any, int]]: - triples = [] - for p in promoters: - cds_same = [c for c in cdss if strand(c) == strand(p)] - if not cds_same: continue - cds_same.sort(key=lambda c: distance(p, c)) - c = cds_same[0] - - term_cands = [t for t in terminators if strand(t) == strand(c)] - # prefer ordered downstream terminators, but allow out-of-order partial credit - term_cands.sort(key=lambda t: distance(c, t)) - t = term_cands[0] if term_cands else None - - pts = 0 - # promoter -> CDS - if in_order_same_strand(p, c): - pts += 5 - pts += prox_points(distance(p, c), 5) - else: - pts += min(3, prox_points(distance(p, c), 5)) # out-of-order partial - - # CDS -> terminator - if t is not None: - if in_order_same_strand(c, t): - pts += 5 - pts += prox_points(distance(c, t), 5) - else: - pts += min(3, prox_points(distance(c, t), 5)) # out-of-order partial - - if role(c) == "marker" and distance(p, c) <= 300: - pts += 2 - - triples.append((p, c, t, pts)) - triples.sort(key=lambda x: x[3], reverse=True) - return triples[:2] - - for (_, _, _, pts) in best_cassettes(): - score += float(min(20, pts)) - - # ---- Standalone promoter & terminator credit ---- - if promoters: - score += float(min(5.0, 1.0 * len(promoters))) # +1 each up to +5 - if terminators: - score += float(min(5.0, 1.0 * len(terminators))) # +1 each up to +5 - - # ---- Payload (GOI) anywhere ---- - target_keywords = [k.lower() for k in (target_keywords or [])] - def is_payload(c) -> bool: - r = role(c) - id_l = (getattr(c, "id", "") or "").lower() - return r in {"payload", "goi", "reporter"} or (target_keywords and any(k in id_l for k in target_keywords)) - - payloads = [c for c in cdss if is_payload(c)] - if payloads: - score += 8.0 - # +4 if any promoter within 500 bp (either direction; same strand preferred implicitly by proximity) - if any(distance(p, c) <= 500 for c in payloads for p in promoters): - score += 4.0 - - # ---- Length penalty: gentler at 3–5 kb ---- - L = max(0, len(sequence or "")) - - def length_penalty(L: int) -> float: - # piecewise linear: - # 1–3 kb: 0 → -3 - # 3–5 kb: -3 → -5 - # 5–7.5 kb:-5 → -9 - # 7.5–10 kb:-9 → -15 - if L <= 1000: - return 0.0 - if L <= 3000: - # span 2kb to -3 - return 3.0 * (L - 1000) / 2000.0 - if L <= 5000: - # next 2kb adds -2 (to -5) - return 3.0 + 2.0 * (L - 3000) / 2000.0 - if L <= 7500: - # next 2.5kb adds -4 (to -9) - return 5.0 + 4.0 * (L - 5000) / 2500.0 - if L <= 10000: - # final 2.5kb adds -6 (to -15) - return 9.0 + 6.0 * (L - 7500) / 2500.0 - return 15.0 - - score -= length_penalty(L) - - # ---- Clamp ---- - return float(max(0.0, min(100.0, score))) - + score01, _ = _SCORER.score(sequence) + return float(max(0.0, min(100.0, 100.0 * score01))) def score_completions(completions: list[str]) -> list[float]: - if _REWARD_LOG_TIMINGS: t0 = time.perf_counter() if not completions: logger.warning("reward.score_completions called with empty completions list") return [] - - # Handle empty sequences efficiently - scores = [] - non_empty_indices = [] - non_empty_completions = [] - + + scores: List[float] = [] + non_empty_indices: List[int] = [] + non_empty_completions: List[str] = [] + for i, c in enumerate(completions): if not c or len(c.strip()) == 0: scores.append(0.0) else: non_empty_indices.append(i) non_empty_completions.append(c) - scores.append(None) # placeholder - - # Annotate and score only non-empty completions + scores.append(0.0) # placeholder + if non_empty_completions: - annotations = annotate_completions(non_empty_completions) - for idx, annotation in zip(non_empty_indices, annotations): - scores[idx] = score_sequence(completions[idx], annotation) - + for idx in non_empty_indices: + scores[idx] = score_sequence(completions[idx]) + if _REWARD_LOG_TIMINGS: dt_ms = (time.perf_counter() - t0) * 1000.0 n = len(scores) mean_score = sum(scores) / n if n else 0.0 - min_score = min(scores) if n else 0.0 - max_score = max(scores) if n else 0.0 logger.info( - f"reward.score n={n} mean={mean_score:.2f} min={min_score:.2f} max={max_score:.2f} time_ms={dt_ms:.2f}" + f"reward.score n={n} mean={mean_score:.2f} time_ms={dt_ms:.2f}" ) return scores diff --git a/src/rewards/scorer.py b/src/rewards/scorer.py index 2d617f4..8aab34f 100644 --- a/src/rewards/scorer.py +++ b/src/rewards/scorer.py @@ -1,26 +1,151 @@ -from rewards.reward_config import RewardConfig +from src.rewards.reward_config import RewardConfig import plasmidkit as pk -import pandas as pd -from typing import Any +from typing import Any, List, Tuple, Union, Dict from concurrent.futures import ThreadPoolExecutor +import time +import hashlib +import os class Scorer: def __init__(self, reward_config: RewardConfig): self.reward_config = reward_config - self.score_functions = [self.score_ori, self.score_promoter, self.score_terminator, self.score_marker, self.score_cds] - self.weights = [self.reward_config.ori_weight, self.reward_config.promoter_weight, self.reward_config.terminator_weight, self.reward_config.marker_weight, self.reward_config.cds_weight] + # Fixed set for simplicity; location/length are folded into component scores/final modulation + self.score_functions = [ + self.score_ori, + self.score_promoter, + self.score_terminator, + self.score_marker, + self.score_cds, + ] + self.weights = [ + self.reward_config.ori_weight, + self.reward_config.promoter_weight, + self.reward_config.terminator_weight, + self.reward_config.marker_weight, + self.reward_config.cds_weight, + ] total_weight = sum(self.weights) - if total_weight: + if total_weight > 0: self.weights = [w / total_weight for w in self.weights] else: - uniform_weight = 1.0 / len(self.score_functions) - self.weights = [uniform_weight for _ in self.score_functions] + uniform = 1.0 / len(self.score_functions) + self.weights = [uniform for _ in self.score_functions] def annotate(self, sequence: str) -> Any: - return pk.annotate(sequence, is_sequence=True) + raw = pk.annotate(sequence, is_sequence=True) + return self._preprocess_annotations(raw) - def runner(self, sequence: str, annotations: Any) -> list[float]: + # --- preprocessing helpers --- + class _Feat: + def __init__(self, type: str, id: str | None, start: int, end: int, strand: str | None, evidence: Any = None): + self.type = type + self.id = id + self.start = int(start) + self.end = int(end) + self.strand = strand or "+" + self.evidence = evidence or {} + + @staticmethod + def _len(x: Any) -> int: + return max(0, int(getattr(x, "end", 0)) - int(getattr(x, "start", 0))) + + @staticmethod + def _overlap_len(a: Any, b: Any) -> int: + s1, e1 = int(getattr(a, "start", 0)), int(getattr(a, "end", 0)) + s2, e2 = int(getattr(b, "start", 0)), int(getattr(b, "end", 0)) + lo = max(min(s1, e1), min(s2, e2)) + hi = min(max(s1, e1), max(s2, e2)) + return max(0, hi - lo) + + def _to_feat(self, x: Any) -> "Scorer._Feat": + return Scorer._Feat( + type=(getattr(x, "type", None) or "").lower(), + id=getattr(x, "id", None), + start=int(getattr(x, "start", 0)), + end=int(getattr(x, "end", 0)), + strand=getattr(x, "strand", "+"), + evidence=getattr(x, "evidence", {}), + ) + + def _merge_group(self, feats: List[Any], threshold: float) -> List["Scorer._Feat"]: + if not feats: + return [] + items = [self._to_feat(f) for f in feats] + items.sort(key=lambda f: (f.strand, f.start, f.end)) + merged: List[Scorer._Feat] = [] + cur = items[0] + for nxt in items[1:]: + ovl = self._overlap_len(cur, nxt) + min_len = max(1, min(self._len(cur), self._len(nxt))) + if ovl / float(min_len) >= threshold and cur.strand == nxt.strand: + cur.start = min(cur.start, nxt.start) + cur.end = max(cur.end, nxt.end) + cur.id = f"{cur.id}|{nxt.id}" if cur.id or nxt.id else None + else: + merged.append(cur) + cur = nxt + merged.append(cur) + return merged + + def _preprocess_annotations(self, annotations: Any) -> List[Any]: + feats = list(annotations) + thr = float(self.reward_config.overlap_merge_threshold) + type_key = lambda x: (getattr(x, "type", None) or "").lower() + + # collect groups + groups: Dict[str, List[Any]] = {} + for f in feats: + groups.setdefault(type_key(f), []).append(f) + + # merge per group for relevant types + merged_groups: Dict[str, List[Scorer._Feat]] = {} + for t in ("rep_origin", "ori", "origin_of_replication", "promoter", "terminator", "marker", "cds"): + if t in groups: + merged_groups[t] = self._merge_group(groups[t], thr) + + # suppress CDS if overlaps any non-CDS (ori/promoter/terminator/marker) + non_cds: List[Scorer._Feat] = [] + for t in ("rep_origin", "ori", "origin_of_replication", "promoter", "terminator", "marker"): + non_cds.extend(merged_groups.get(t, [])) + + filtered_cds: List[Scorer._Feat] = [] + for c in merged_groups.get("cds", []): + if any(self._overlap_len(c, o) > 0 for o in non_cds): + continue + filtered_cds.append(c) + merged_groups["cds"] = filtered_cds + + # rebuild final list: prefer merged groups for those types; keep other features as-is + final: List[Any] = [] + merged_types = set(merged_groups.keys()) + for t, items in merged_groups.items(): + final.extend(items) + for f in feats: + t = type_key(f) + if t not in merged_types: + final.append(f) + return final + + @staticmethod + def _read_fasta_file(path: str) -> str: + with open(path, "r") as f: + lines = [ln.strip() for ln in f.readlines()] + seq_lines = [ln for ln in lines if ln and not ln.startswith(">")] + return "".join(seq_lines).replace(" ", "").upper() + @staticmethod + def _derive_source(sequence: str, provided: str | None) -> str: + if provided: + return provided + # Try FASTA header if present + if sequence and sequence.lstrip().startswith(">"): + first = sequence.splitlines()[0].lstrip()[1:].strip() + return first.split()[0] if first else "" + # Fallback: length + sha1 fingerprint + sha = hashlib.sha1((sequence or "").encode("utf-8")).hexdigest()[:8] + return f"seq:bp{len(sequence or '')}:sha{sha}" + + def runner(self, sequence: str, annotations: Any) -> List[float]: """ Takes a list of runnable functions with the same signature and runs them in parallel. @@ -35,24 +160,186 @@ def runner(self, sequence: str, annotations: Any) -> list[float]: futures = [executor.submit(func, sequence, annotations) for func in self.score_functions] return [future.result() for future in futures] - def score_ori(self, seq: str, annotations: pd.DataFrame) -> float: - return 0.0 + # --- helpers --- + @staticmethod + def _feat_type(x: Any) -> str: + return (getattr(x, "type", None) or "").lower() + + @staticmethod + def _feat_id(x: Any) -> str: + return (getattr(x, "id", None) or "").lower() - def score_promoter(self, seq: str, annotations: pd.DataFrame) -> float: - return 0.0 + @staticmethod + def _strand(x: Any) -> str: + return getattr(x, "strand", "+") - def score_terminator(self, seq: str, annotations: pd.DataFrame) -> float: - return 0.0 + @staticmethod + def _start(x: Any) -> int: + return int(getattr(x, "start", 0)) - def score_marker(self, seq: str, annotations: pd.DataFrame) -> float: - return 0.0 + @staticmethod + def _end(x: Any) -> int: + return int(getattr(x, "end", 0)) - def score_cds(self, seq: str, annotations: pd.DataFrame) -> float: - return 0.0 + @staticmethod + def _distance(a: Any, b: Any) -> int: + if Scorer._end(a) < Scorer._start(b): + return Scorer._start(b) - Scorer._end(a) + if Scorer._end(b) < Scorer._start(a): + return Scorer._start(a) - Scorer._end(b) + return 0 - def score_length(self, seq: str, annotations: pd.DataFrame) -> float: - return 0.0 + @staticmethod + def _in_order_same_strand(a: Any, b: Any) -> bool: + if Scorer._strand(a) != Scorer._strand(b): + return False + if Scorer._strand(a) == "+": + return Scorer._start(a) <= Scorer._start(b) + return Scorer._end(a) >= Scorer._end(b) - def score(self, seq: str, annotations: pd.DataFrame) -> float: + def _prox_points(self, d: int) -> float: + return ( + float(self.reward_config.cassette_proximity_points) + if d <= int(self.reward_config.proximity_threshold_bp) + else 0.0 + ) + + @staticmethod + def _filter_allowed(xs: List[Any], allowed: List[str] | None) -> List[Any]: + if not allowed: + return xs + allowed_l = [a.lower() for a in allowed] + return [x for x in xs if any(a in Scorer._feat_id(x) for a in allowed_l)] + + def _count_score(self, count: int, min_req: int, max_allowed: int) -> float: + if max_allowed <= 0: + return 0.0 + # Below minimum: proportion of requirement; softer if punish_mode + if count < min_req: + proportion = (count / float(min_req)) if min_req > 0 else 0.0 + return max(0.0, min(1.0, proportion * (0.5 if self.reward_config.punish_mode else 1.0))) + # Above maximum: penalize if punishing, otherwise cap at full credit + if count > max_allowed: + return float(self.reward_config.violation_penalty_factor) if self.reward_config.punish_mode else 1.0 + # Within [min, max]: full credit + return 1.0 + + def _compute_cassette_points(self, promoters: List[Any], cdss: List[Any], terminators: List[Any]) -> float: + # Returns summed raw points across top-N cassettes + triples: List[Tuple[Any, Any, Any, float]] = [] + for p in promoters: + cds_same = [c for c in cdss if self._strand(c) == self._strand(p)] + if not cds_same: + continue + cds_same.sort(key=lambda c: self._distance(p, c)) + c = cds_same[0] + term_cands = [t for t in terminators if self._strand(t) == self._strand(c)] + term_cands.sort(key=lambda t: self._distance(c, t)) + t = term_cands[0] if term_cands else None + pts = 0.0 + # promoter -> CDS + if self._in_order_same_strand(p, c): + pts += self.reward_config.cassette_order_points + pts += min(self.reward_config.cassette_proximity_points, self._prox_points(self._distance(p, c))) + else: + pts += min(3.0, self._prox_points(self._distance(p, c))) + # CDS -> terminator + if t is not None: + if self._in_order_same_strand(c, t): + pts += self.reward_config.cassette_order_points + pts += min(self.reward_config.cassette_proximity_points, self._prox_points(self._distance(c, t))) + else: + pts += min(3.0, self._prox_points(self._distance(c, t))) + triples.append((p, c, t, pts)) + triples.sort(key=lambda x: x[3], reverse=True) + top = triples[: int(self.reward_config.cassette_max_cassettes)] + total = 0.0 + for (_, _, _, pts) in top: + total += min(self.reward_config.cassette_max_points_per, float(pts)) + return total + + # --- component scores (0..1) --- + def score_ori(self, seq: str, annotations: Any) -> float: + feats = list(annotations) + oris = [x for x in feats if self._feat_type(x) in ("rep_origin", "ori", "origin_of_replication")] + oris = self._filter_allowed(oris, self.reward_config.allowed_oris) + return self._count_score(len(oris), self.reward_config.ori_min, self.reward_config.ori_max) + + def score_promoter(self, seq: str, annotations: Any) -> float: + feats = list(annotations) + promoters = [x for x in feats if self._feat_type(x) == "promoter"] + promoters = self._filter_allowed(promoters, self.reward_config.allowed_promoters) + score = self._count_score(len(promoters), self.reward_config.promoter_min, self.reward_config.promoter_max) + # small standalone credit (capped implicitly by count scoring) + return score + + def score_terminator(self, seq: str, annotations: Any) -> float: + feats = list(annotations) + terms = [x for x in feats if self._feat_type(x) == "terminator"] + terms = self._filter_allowed(terms, self.reward_config.allowed_terminators) + return self._count_score(len(terms), self.reward_config.terminator_min, self.reward_config.terminator_max) + + def score_marker(self, seq: str, annotations: Any) -> float: + feats = list(annotations) + markers = [x for x in feats if self._feat_type(x) == "marker"] + markers = self._filter_allowed(markers, self.reward_config.allowed_markers) + if markers: + return 1.0 + return 0.0 if self.reward_config.punish_mode else 0.5 + + def score_cds(self, seq: str, annotations: Any) -> float: + feats = list(annotations) + cdss = [x for x in feats if self._feat_type(x) == "cds"] + cdss = self._filter_allowed(cdss, self.reward_config.allowed_cds) + base = self._count_score(len(cdss), self.reward_config.cds_min, self.reward_config.cds_max) + + if not self.reward_config.location_aware: + return base + + promoters = [x for x in feats if self._feat_type(x) == "promoter"] + terminators = [x for x in feats if self._feat_type(x) == "terminator"] + cassette_pts = self._compute_cassette_points(promoters, cdss, terminators) + max_total = self.reward_config.cassette_max_cassettes * self.reward_config.cassette_max_points_per + bonus_scale = float(self.reward_config.location_bonus_scale) + cassette_bonus = bonus_scale * (cassette_pts / max_total if max_total > 0 else 0.0) + return max(0.0, min(1.0, base + cassette_bonus)) + + def _length_factor(self, seq: str) -> float: + if not self.reward_config.length_penalty: + return 1.0 + L = len(seq or "") + mn = self.reward_config.min_length + mx = self.reward_config.max_length + if mn is None and mx is None: + return 1.0 + # inside range: full credit; outside: halve if punish_mode else no change + in_range = (mn is None or L >= mn) and (mx is None or L <= mx) + if in_range: + return 1.0 + return float(self.reward_config.violation_penalty_factor) if self.reward_config.punish_mode else 1.0 + + def score(self, seq: str, source: str | None = None) -> Tuple[float, Dict[str, float]]: + t0 = time.perf_counter() + src = Scorer._derive_source(seq, source) + annotations = self.annotate(seq) results = self.runner(seq, annotations) - return sum(w * r for w, r in zip(self.weights, results)) \ No newline at end of file + ori, prom, term, mark, cds = results + base = sum(w * r for w, r in zip(self.weights, results)) + length_factor = self._length_factor(seq) + final = max(0.0, min(1.0, base * length_factor)) + components: Dict[str, float] = { + "ori": float(ori), + "promoter": float(prom), + "terminator": float(term), + "marker": float(mark), + "cds": float(cds), + "length_factor": float(length_factor), + } + dt_ms = (time.perf_counter() - t0) * 1000.0 + cfg = self.reward_config.model_dump(exclude_none=True) + print(f"source={src} config={cfg} score={final:.4f} len={len(seq)} time_ms={dt_ms:.2f} components={components}") + return float(final), components + + def score_fasta(self, fasta_path: str) -> Tuple[float, Dict[str, float]]: + seq = Scorer._read_fasta_file(fasta_path) + return self.score(seq, source=os.path.basename(fasta_path)) \ No newline at end of file diff --git a/tests/data/RF0G-IodoY.fasta b/tests/data/RF0G-IodoY.fasta new file mode 100644 index 0000000..1eca288 --- /dev/null +++ b/tests/data/RF0G-IodoY.fasta @@ -0,0 +1,79 @@ +> pRF0G-IodoY sequence 5321 bps +CTGAGCAAACTGGCCTCAGGCATTTGAGAAGCACACGGTCACACTGCTTCCGGTAGTCAATAAACCGGTA +AACCAGCAATAGACATAAGCGGCTATTTAACGACCCTGCCCTGAACCGACGACCGGGTCGAATTTGCTTT +CGAATTTCTGCCATTCATCCGCTTATTATCACTTATTCAGGCGTAGCACCAGGCGTTTAAGGGCACCAAT +AACTGCCTTAAAAAAATTACGCCCCGCCCTGCCACTCATCGCAGTACTTAGGTGGCGGTACTTGGGTCGA +TATCAAAGTGCATCACTTCTTCCCGTATGCCCAACTTTGTATAGAGAGCCACTGCGGGATCGTCACCGTA +ATCTGCTTGCACGTAGATCACATAAGCACCAAGCGCGTTGGCCTCATGCTTGAGGAGATTGATGAGCGCG +GTGGCAATGCCCTGCCTCCGGTGCTCGCCGGAGACTGCGAGATCATAGATATAGATCTCACTACGCGGCT +GCTCAAACTTGGGCAGAACGTAAGCCGCGAGAGCGCCAACAACCGCTTCTTGGTCGAAGGCAGCAAGCGC +GATGAATGTCTTACTACGGAGCAAGTTCCCGAGGTAATCGGAGTCCGGCTGATGTTGGGAGTAGGTGGCT +ACGTCTCCGAACTCACGACCGAAAAGATCAAGAGCAGCCCGCATGGATTTGACTTGGTCAGGGCCGAGCC +TACATGTGCGAATGATGCCCATACTTGAGCCACCTAACTTTGTTTTAGGGCGACTGCCCTGCTGCGTAAC +ATCGTTGCTGCTGCGTAACATAACACCCCTTGTATTACTGTTTATGTAAGCAGACAGTTTTATTGTTCAT +GATGATATATTTTTATCTTGTGCAATGTAACATCAGAGATTTTGAGACACAACGTGGCTTTCCGAATTCC +GGATGAGCATTCATCAGGCGGGCAAGAATGTGAATAAAGGCCGGATAAAACTTGTGCTTATTTTTCTTTA +CGGTCTTTAAAAAGGCCGTAATATCCAGCTGAACGGTCTGGTTATAGGTACATTGAGCAACTGACTGAAA +TGCCTCAAAATGTTCTTTACGATGCCATTGGGATATATCAACGGTGGTATATCCAGTGATTTTTTTCTCC +ATTTTAGCTTCCTTAGCTCCTGAAAATCTCGATAACTCAAAAAATACGCCCGGTAGTGATCTTATTTCAT +TATGGTGAAAGTTGGAACCTCTTACGTGCCGATCAACGTCTCATTTTCGCCAAAAGTTGGCCCAGGGCTT +CCCGGTATCAACAGGGACACCAGGATTTATTTATTCTGCGAAGTGATCTTCCGTCACAGGTATTTATTCG +GCGCAAAGTGCGTCGGGTGATGCTGCCAACTTACTGATTTAGTGTATGATGGTGTTTTTGAGGTGCTCCA +GTGGCTTCTGTTTCTATCAGCTGTCCCTCCTGTTCAGCTACTGACGGGGTGGTGCGTAACGGCAAAAGCA +CCGCCGGACATCAGCGCTAGCGGAGTGTATACTGGCTTACTATGTTGGCACTGATGAGGGTGTCAGTGAA +GTGCTTCATGTGGCAGGAGAAAAAAGGCTGCACCGGTGCGTCAGCAGAATATGTGATACAGGATATATTC +CGCTTCCTCGCTCACTGACTCGCTACGCTCGGTCGTTCGACTGCGGCGAGCGGAAATGGCTTACGAACGG +GGCGGAGATTTCCTGGAAGATGCCAGGAAGATACTTAACAGGGAAGTGAGAGGGCCGCGGCAAAGCCGTT +TTTCCATAGGCTCCGCCCCCCTGACAAGCATCACGAAATCTGACGCTCAAATCAGTGGTGGCGAAACCCG +ACAGGACTATAAAGATACCAGGCGTTTCCCCCTGGCGGCTCCCTCGTGCGCTCTCCTGTTCCTGCCTTTC +GGTTTACCGGTGTCATTCCGCTGTTATGGCCGCGTTTGTCTCATTCCACGCCTGACACTCAGTTCCGGGT +AGGCAGTTCGCTCCAAGCTGGACTGTATGCACGAACCCCCCGTTCAGTCCGACCGCTGCGCCTTATCCGG +TAACTATCGTCTTGAGTCCAACCCGGAAAGACATGCAAAAGCACCACTGGCAGCAGCCACTGGTAATTGA +TTTAGAGGAGTTAGTCTTGAAGTCATGCGCCGGTTAAGGCTAAACTGAAAGGACAAGTTTTGGTGACTGC +GCTCCTCCAAGCCAGTTACCTCGGTTCAAAGAGTTGGTAGCTCAGAGAACCTTCGAAAAACCGCCCTGCA +AGGCGGTTTTTTCGTTTTCAGAGCAAGAGATTACGCGCAGACCAAAACGATCTCAAGAAGATCATCTTAT +TAATCAGATAAAATATTTCTAGATTTCAGTGCAATTTATCTCTTCAAATGTAGCACCTGAAGTCAGCCCC +ATACGATATAAGTTGTAATTCTCATGTTTGACAGCATTATCATCGATAAGCTTGATGCGTGGAAGATTGA +TCGTCTTGCACCCTGAAAAGATGCAAAAATCTTGCTTTAATCGCTGGTACTCCTGATTCTGGCACTTTAT +TCTATGTCTCTTTCGCATCTGGCGAAAAGTCGTGTACCGGCAAAGGTGCAGTCGTTATATACTTGGAGAT +TCATATGGACGAGTTCGAGATGATCAAACGCAACACCAGCGAAATCATCAGCGAAGAAGAACTGCGCGAA +GTGCTGAAAAAAGACGAGAAGAGCGCCTATATTGGCTTCGAACCGAGCGGCAAAATTCATCTGGGCCACT +ACCTGCAGATCAAGAAGATGATCGATCTGCAGAACGCGGGCTTCGATATCATCATTCTGCTGGCCGATCT +GGCCGCATACCTGAACCAGAAAGGCGAACTGGACGAAATCCGCAAAATCGGCGACTACAACAAGAAAGTG +TTCGAAGCGATGGGCCTGAAAGCGAAATACGTCTATGGCAGCGAATTTCAGCTGGACAAAGACTACACCC +TGAACGTGTATCGTCTGGCACTGAAAACCACCCTGAAACGTGCTCGTCGTTCTATGGAACTGATTGCCCG +CGAAGATGAAAACCCGAAAGTGGCGGAAGTGATCTATCCGATCATGCAGGTGAACACCAGCCACTATCTG +GGCGTAGATGTAGCAGTTGGTGGCATGGAACAGCGCAAAATCCACATGCTGGCTCGTGAACTGCTGCCGA +AAAAAGTCGTGTGCATCCACAACCCGGTTTTAACGGGTCTGGATGGTGAAGGCAAAATGAGCAGCAGCAA +AGGCAACTTTATCGCCGTGGATGATTCTCCGGAAGAAATCCGCGCGAAGATCAAAAAAGCGTATTGCCCG +GCAGGTGTAGTTGAAGGTAACCCGATCATGGAGATCGCGAAATACTTCCTGGAATACCCGCTGACCATTA +AACGCCCGGAAAAATTCGGTGGTGATCTGACCGTGAACAGCTACGAAGAACTGGAGAGCCTGTTCAAGAA +CAAAGAACTGCACCCGATGCGTCTGAAAAATGCCGTTGCGGAAGAGCTGATCAAAATCCTGGAACCGATT +CGCAAACGTCTGTAAGATCTGGATCCTCTACGCCGGACGCATCGTGGCCGGCATCACCGGCGCCACAGGT +GCGGTTGCTGGCGCCTATATCGCCGACATCACCGATGGGGAAGATCGGGCTCGCCACTTCGGGCTCATGA +GGCATGCGGCGCCGCTTCTTTGAGCGAACGATCAAAAATAAGTGGCGCCCCATCAAAAAAATATTCTCAA +CATAAAAAACTTTGTGTAATACTTGTAACGCTGCCATCAGACGCATTCCGGCGGTAGTTCAGCAGGGCAG +AACGGCGGACTCTAAATCCGCAGGTCGCTGGTTCAAATCCGGCCCGCCGGACCATTTATCACAGATTGGA +AATTTTTGATCCTTAGCGAAAGCTAAGGATTTTTTTTAGTCGACCGATGCCCTTGAGAGCCTTCAACCCA +GTCAGCTCCTTCCGGTGGGCGCGGGGCATGACTATCGTCGCCGCACTTATGACTGTCTTCTTTATCATGC +AACTCGTAGGACAGGTGCCGGCAGCGCTCTGGGTCATTTTCGGCGAGGACCGCTTTCGCTGGAGCGCGAC +GATGATCGGCCTGTCGCTTGCGGTATTCGGAATCTTGCACGCCCTCGCTCAAGCCTTCGTCACTGGTCCC +GCCACCAAACGTTTCGGCGAGAAGCAGGCCATTATCGCCGGCATGGCGGCCGACGCGCTGGGCTACGTCT +TGCTGGCGTTCGCGACGCGAGGCTGGATGGCCTTCCCCATTATGATTCTTCTCGCTTCCGGCGGCATCGG +GATGCCCGCGTTGCAGGCCATGCTGTCCAGGCAGGTAGATGACGACCATCAGGGACAGCTTCAAGGATCG +CTCGCGGCTCTTACCAGCCTAACTTCGATCACTGGACCGCTGATCGTCACGGCGATTTATGCCGCCTCGG +CGAGCACATGGAACGGGTTGGCATGGATTGTAGGCGCCGCCCTATACCTTGTCTGCCTCCCCGCGTTGCG +TCGCGGTGCATGGAGCCGGGCCACCTCGACCTGAATGGAAGCCGGCGGCACCTCGCTAACGGATTCACCA +CTCCAAGAATTGGAGCCAATCAATTCTTGCGGAGAACTGTGAATGCGCAAACCAACCCTTGGCAGAACAT +ATCCATCGCGTCCGCCATCTCCAGCAGCCGCACGCGGCGCATCTCGGGCAGCGTTGGGTCCTGGCCACGG +GTGCGCATGATCGTGCTCCTGTCGTTGAGGACCCGGCTAGGCTGGCGGGGTTGCCTTACTGGTTAGCAGA +ATGAATCACCGATACGCGAGCGAACGTGAAGCGACTGCTGCTGCAAAACGTCTGCGACCTGAGCAACAAC +ATGAATGGTCTTCGGTTTCCGTGTTTCGTAAAGTCTGGAAACGCGGAAGTCCCCTACGTGCTGCTGAAGT +TGCCCGCAACAGAGAGTGGAACCAACCGGTGATACCACGATACTATGACTGAGAGTCAACGCCATGAGCG +GCCTCATTTCTTATTCTGAGTTACAACAGTCCGCACCGCTGTCCGGTAGCTCCTTCCGGTGGGCGCGGGG +CATGACTATCGTCGCCGCACTTATGACTGTCTTCTTTATCATGCAACTCGTAGGACAGGTGCCGGCAGCG +CCCAACAGTCCCCCGGCCACGGGGCCTGCCACCATACCCACGCCGAAACAAGCGCCCTGCACCATTATGT +TCCGGATCTGCATCGCAGGATGCTGCTGGCTACCCTGTGGAACACCTACATCTGTATTAACGAAGCGCTA +ACCGTTTTTATCAGGCTCTGGGAGGCAGAATAAATGATCATATCGTCAATTATTACCTCCACGGGGAGAG +C + \ No newline at end of file diff --git a/tests/data/pSC101.fasta b/tests/data/pSC101.fasta new file mode 100644 index 0000000..9b23642 --- /dev/null +++ b/tests/data/pSC101.fasta @@ -0,0 +1,2 @@ +>Addgene_NGS_Result +GGTCTGCTATGTGGTGCTATCTGACTTTTTGCTGTTCAGCAGTTCCTGCCCTCTGATTTTCCAGTCTGACCACTTCGGATTATCCCGTGACAGGTCATTCAGACTGGCTAATGCACCCAGTAAGGCAGCGGTATCATCAACGGGGTCTGACGCTCAGTGGAACGAAAACTCACGTTAAGGGATTTTGGTCATGAGATTATCAAAAAGGATCTTCACCTAGATCCGTTATGCAGCGGAAAGTAAAAAATTTTTAGTTTATTAGACATCTCCACAAAAGGCGTAGTGTACAGTGACAAATTATCTGTCGTCGGTGACAGATTAATGTCATTGTGACTATTTAATTGTCGTCGTGACCCATCAGCGTTGCTTAATTAATTGATGACAAATTAAATGTCATCAATATAATATGCTCTGCAATTATTATACAAAGCAATTAAAACAAGCGGATAAAAGGACTTGCTTTCAACCCACCCCTAAGTTTAATAGTTACTGAGGGGGATCCACTAGTGAGCTCATGCATGATCTCGAATTAGCTTCAAAAGCGCTCTGAAGTTCCTATACTTTCTAGAGAATAGGAACTTCGGAATAGGAACTTCAAGATCCCCTGATTCCCTTTGTCAACAGCAATGGATAATTCGATTTAACAAATGCATGGCGCAAGGGCTGCTAAAGGAAGCGGAACACGTAGAAAGCCAGTCCGCAGAAACGGTGCTGACCCCGGATGAATGTCAGCTACTGGGCTATCTGGACAAGGGAAAACGCAAGCGCAAAGAGAAAGCAGGTAGCTTGCAGTGGGCTTACATGGCGATAGCTAGACTGGGCGGTTTTATGGACAGCAAGCGAACCGGAATTGCCAGCTGGGGCGCCCTCTGGTAAGGTTGGGAAGCCCTGCAAAGTAAACTGGATGGCTTTCTTGCCGCCAAGGATCTGATGGCGCAGGGGATCAAGATCTGATCAAGAGACAGGATGAGGATCGTTTCGCATGATTGAACAAGATGGATTGCACGCAGGTTCTCCGGCCGCTTGGGTGGAGAGGCTATTCGGCTATGACTGGGCACAACAGACAATCGGCTGCTCTGATGCCGCCGTGTTCCGGCTGTCAGCGCAGGGGCGCCCGGTTCTTTTTGTCAAGACCGACCTGTCCGGTGCCCTGAATGAACTGCAGGACGAGGCAGCGCGGCTATCGTGGCTGGCCACGACGGGCGTTCCTTGCGCAGCTGTGCTCGACGTTGTCACTGAAGCGGGAAGGGACTGGCTGCTATTGGGCGAAGTGCCGGGGCAGGATCTCCTGTCATCCCACCTTGCTCCTGCCGAGAAAGTATCCATCATGGCTGATGCAATGCGGCGGCTGCATACGCTTGATCCGGCTACCTGCCCATTCGACCACCAAGCGAAACATCGCATCGAGCGAGCACGTACTCGGATGGAAGCCGGTCTTGTCGATCAGGATGATCTGGACGAAGAGCATCAGGGGCTCGCGCCAGCCGAACTGTTCGCCAGGCTCAAGGCGCGCATGCCCGACGGCGAGGATCTCGTCGTGACCCATGGCGATGCCTGCTTGCCGAATATCATGGTGGAAAATGGCCGCTTTTCTGGATTCATCGACTGTGGCCGGCTGGGTGTGGCGGACCGCTATCAGGACATAGCGTTGGCTACCCGTGATATTGCTGAAGAGCTTGGCGGCGAATGGGCTGACCGCTTCCTCGTGCTTTACGGTATCGCCGCTCCCGATTCGCAGCGCATCGCCTTCTATCGCCTTCTTGACGAGTTCTTCTGAATTGAAAAAGGAAGAGTATGAGGATCCAACATTTCCAATCACTAGTGAATTATCTAGAATTATTCCATTGAGTAAGTTTTTAAGCACATCAGCTTCAAAAGCGCTCTGAAGTTCCTATACTTTCTAGAGAATAGGAACTTCGGAATAGGTACTTCAAGATCCCCAATTCGAGATCGTCCGGGCCGCAAGCTCCTAGCGGCGGATTTGTCCTACTCAGGAGAGCGTTCACCGACAAACAACAGATAAAACGAAAGGCCCAGTCTTTCGACTGAGCCTTTCGTTTTATTTGATGCCTCAAGCTAGAGAGTCATTACCCCAGGCGTTTAAGGGCACCAATAACTGCCTTAAAAAAATTACGCCCCGCCCTGCCACTCATCGCAGTCTAGCTTGGATTCTCACCAATAAAAAACGCCCGGCGGCAACCGAGCGTTCTGAACAAATCCAGATGGAGTTCTGAGGTCATTACTGGATCTATCAACAGGAGTCCAAGCTCAGCTAATTAAGGCGACAGTCAATTTGTCATTATGAAAATACACAAAAGCTTTTTCCTATCTTGCAAAGCGACAGCTAATTTGTCACAATCACGGACAACGACATCTATTTTGTCACTGCAAAGAGGTTATGCTAAAACTGCCAAAGCGCTATAATCTATACTGTATAAGGATTTTACTGATGACAATAATTTGTCACAACGACATATAATTAGTCACTGTACACGTAGAGACGTAGCAATGCTACCTCTCTACAATGGTTTTGTGTTAGTCTTGATGCTTCACTGATAGATACAAGAGCCATAAGAACCTCAGATCCTTCCGTATTTAGCCAGTATGTTCTCTAGTGTGGTTCGTTGTTTTTGCGTGAGCCATGAGAACGAACCATTGAGATCATACTTACTTTGCATGTCACTCAAAAATTTTGCCTCAAAACTGGTGAGCTGAATTTTTGCAGTTAAAGCATCGTGTAGTGTTTTTCTTAGTCCGTTATGTAGGTAGGAATCTGATGTAATGGTTGTTGGTATTTTGTCACCATTCATTTTTATCTGGTTGTTCTCAAGTTCGGTTACGAGATCCATTTGTCTATCTAGTTCAACTTGGAAAATCAACGTATCAGTCGGGCGGCCTCGCTTATCAACCACCAATTTCATATTGCTGTAAGTGTTTAAATCTTTACTTATTGGTTTCAAAACCCATTGGTTAAGCCTTTTAAACTCATGGTAGTTATTTTCAAGCATTAACATGAACTTAAATTCATCAAGGCTAATCTCTATATTTGCCTTGTGAGTTTTCTTTTGTGTTAGTTCTTTTAATAACCACTCATAAATCCTCATAGAGTATTTGTTTTCAAAAGACTTAACATGTTCCAGATTATATTTTATGAATTTTTTTAACTGGAAAAGATAAGGCAATATCTCTTCACTAAAAACTAATTCTAATTTTTCGCTTGAGAACTTGGCATAGTTTGTCCACTGGAAAATCTCAAAGCCTTTAACCAAAGGATTCCTGATTTCCACAGTTCTCGTCATCAGCTCTCTGGTTGCTTTAGCTAATACACCATAAGCATTTTCCCTACTGATGTTCATCATCTGAACGTATTGGTTATAAGTGAACGATACCGTCCGTTCTTTCCTTGTAGGGTTTTCAATCGTGGGGTTGAGTAGTGCCACACAGCATAAAATTAGCTTGGTTTCATGCTCCGTTAAGTCATAGCGACTAATCGCTAGTTCATTTGCTTTGAAAACAACTAATTCAGACATACATCTCAATTGGTCTAGGTGATTTTAATCACTATACCAATTGAGATGGGCTAGTCAATGATAATTACTAGTCCTTTTCCTTTGAGTTGTGGGTATCTGTAAATTCTGCTAGACCTTTGCTGGAAAACTTGTAAATTCTGCTAGACCCTCTGTAAATTCCGCTAGACCTTTGTGTGTTTTTTTTGTTTATATTCAAGTGGTTATAATTTATAGAATAAAGAAAGAATAAAAAAAGATAAAAAGAATAGATCCCAGCCCTGTGTATAACTCACTACTTTAGTCAGTTCCGCAGTATTACAAAAGGATGTCGCAAACGCTGTTTGCTCCTCTACAAAACAGACCTTAAAACCCTAAAGGCTTAAGTAGCACCCTCGCAAGCTCGGTTGCGGCCGCAATCGGGCAAATCGCTGAATATTCCTTTTGTCTCCGACCATCAGGCACCTGAGTCGCTGTCTTTTTCGTGACATTCAGTTCGCTGCGCTCACGGCTCTGGCAGTGAATGGGGGTAAATGGCACTACAGGCGCCTTTTATGGATTCATGCAAGGAAACTACCCATAATACAAGAAAAGCCCGTCACGGGCTTCTCAGGGCGTTTTATGGCG \ No newline at end of file diff --git a/tests/data/pUC19.fasta b/tests/data/pUC19.fasta new file mode 100644 index 0000000..489306c --- /dev/null +++ b/tests/data/pUC19.fasta @@ -0,0 +1,2 @@ +>Addgene_NGS_Result +GAGATACCTACAGCGTGAGCTATGAGAAAGCGCCACGCTTCCCGAAGGGAGAAAGGCGGACAGGTATCCGGTAAGCGGCAGGGTCGGAACAGGAGAGCGCACGAGGGAGCTTCCAGGGGGAAACGCCTGGTATCTTTATAGTCCTGTCGGGTTTCGCCACCTCTGACTTGAGCGTCGATTTTTGTGATGCTCGTCAGGGGGGCGGAGCCTATGGAAAAACGCCAGCAACGCGGCCTTTTTACGGTTCCTGGCCTTTTGCTGGCCTTTTGCTCACATGTTCTTTCCTGCGTTATCCCCTGATTCTGTGGATAACCGTATTACCGCCTTTGAGTGAGCTGATACCGCTCGCCGCAGCCGAACGACCGAGCGCAGCGAGTCAGTGAGCGAGGAAGCGGAAGAGCGCCCAATACGCAAACCGCCTCTCCCCGCGCGTTGGCCGATTCATTAATGCAGCTGGCACGACAGGTTTCCCGACTGGAAAGCGGGCAGTGAGCGCAACGCAATTAATGTGAGTTAGCTCACTCATTAGGCACCCCAGGCTTTACACTTTATGCTTCCGGCTCGTATGTTGTGTGGAATTGTGAGCGGATAACAATTTCACACAGGAAACAGCTATGACCATGATTACGCCAAGCTTGCATGCCTGCAGGTCGACTCTAGAGGATCCCCGGGTACCGAGCTCGAATTCACTGGCCGTCGTTTTACAACGTCGTGACTGGGAAAACCCTGGCGTTACCCAACTTAATCGCCTTGCAGCACATCCCCCTTTCGCCAGCTGGCGTAATAGCGAAGAGGCCCGCACCGATCGCCCTTCCCAACAGTTGCGCAGCCTGAATGGCGAATGGCGCCTGATGCGGTATTTTCTCCTTACGCATCTGTGCGGTATTTCACACCGCATATGGTGCACTCTCAGTACAATCTGCTCTGATGCCGCATAGTTAAGCCAGCCCCGACACCCGCCAACACCCGCTGACGCGCCCTGACGGGCTTGTCTGCTCCCGGCATCCGCTTACAGACAAGCTGTGACCGTCTCCGGGAGCTGCATGTGTCAGAGGTTTTCACCGTCATCACCGAAACGCGCGAGACGAAAGGGCCTCGTGATACGCCTATTTTTATAGGTTAATGTCATGATAATAATGGTTTCTTAGACGTCAGGTGGCACTTTTCGGGGAAATGTGCGCGGAACCCCTATTTGTTTATTTTTCTAAATACATTCAAATATGTATCCGCTCATGAGACAATAACCCTGATAAATGCTTCAATAATATTGAAAAAGGAAGAGTATGAGTATTCAACATTTCCGTGTCGCCCTTATTCCCTTTTTTGCGGCATTTTGCCTTCCTGTTTTTGCTCACCCAGAAACGCTGGTGAAAGTAAAAGATGCTGAAGATCAGTTGGGTGCACGAGTGGGTTACATCGAACTGGATCTCAACAGCGGTAAGATCCTTGAGAGTTTTCGCCCCGAAGAACGTTTTCCAATGATGAGCACTTTTAAAGTTCTGCTATGTGGCGCGGTATTATCCCGTATTGACGCCGGGCAAGAGCAACTCGGTCGCCGCATACACTATTCTCAGAATGACTTGGTTGAGTACTCACCAGTCACAGAAAAGCATCTTACGGATGGCATGACAGTAAGAGAATTATGCAGTGCTGCCATAACCATGAGTGATAACACTGCGGCCAACTTACTTCTGACAACGATCGGAGGACCGAAGGAGCTAACCGCTTTTTTGCACAACATGGGGGATCATGTAACTCGCCTTGATCGTTGGGAACCGGAGCTGAATGAAGCCATACCAAACGACGAGCGTGACACCACGATGCCTGTAGCAATGGCAACAACGTTGCGCAAACTATTAACTGGCGAACTACTTACTCTAGCTTCCCGGCAACAATTAATAGACTGGATGGAGGCGGATAAAGTTGCAGGACCACTTCTGCGCTCGGCCCTTCCGGCTGGCTGGTTTATTGCTGATAAATCTGGAGCCGGTGAGCGTGGGTCTCGCGGTATCATTGCAGCACTGGGGCCAGATGGTAAGCCCTCCCGTATCGTAGTTATCTACACGACGGGGAGTCAGGCAACTATGGATGAACGAAATAGACAGATCGCTGAGATAGGTGCCTCACTGATTAAGCATTGGTAACTGTCAGACCAAGTTTACTCATATATACTTTAGATTGATTTAAAACTTCATTTTTAATTTAAAAGGATCTAGGTGAAGATCCTTTTTGATAATCTCATGACCAAAATCCCTTAACGTGAGTTTTCGTTCCACTGAGCGTCAGACCCCGTAGAAAAGATCAAAGGATCTTCTTGAGATCCTTTTTTTCTGCGCGTAATCTGCTGCTTGCAAACAAAAAAACCACCGCTACCAGCGGTGGTTTGTTTGCCGGATCAAGAGCTACCAACTCTTTTTCCGAAGGTAACTGGCTTCAGCAGAGCGCAGATACCAAATACTGTTCTTCTAGTGTAGCCGTAGTTAGGCCACCACTTCAAGAACTCTGTAGCACCGCCTACATACCTCGCTCTGCTAATCCTGTTACCAGTGGCTGCTGCCAGTGGCGATAAGTCGTGTCTTACCGGGTTGGACTCAAGACGATAGTTACCGGATAAGGCGCAGCGGTCGGGCTGAACGGGGGGTTCGTGCACACAGCCCAGCTTGGAGCGAACGACCTACACCGAACT \ No newline at end of file diff --git a/tests/test_rewards.py b/tests/test_rewards.py new file mode 100644 index 0000000..e250494 --- /dev/null +++ b/tests/test_rewards.py @@ -0,0 +1,98 @@ +import os +import pytest + +from src.rewards.scorer import Scorer +from src.rewards.reward_config import RewardConfig +from src.rewards import rewards as rewards_mod + + +def _read_fasta_sequence(path: str) -> str: + with open(path, "r") as f: + lines = [line.strip() for line in f.readlines()] + seq_lines = [ln for ln in lines if not ln.startswith(">") and ln] + return "".join(seq_lines).replace(" ", "").replace("\n", "").upper() + + +def test_scorer_components_and_bounds(): + data_dir = os.path.join(os.path.dirname(__file__), "data") + fasta_paths = [ + os.path.join(data_dir, "pUC19.fasta"), + os.path.join(data_dir, "pSC101.fasta"), + os.path.join(data_dir, "RF0G-IodoY.fasta"), + ] + cfg = RewardConfig() + scorer = Scorer(cfg) + + for p in fasta_paths: + score, components = scorer.score_fasta(p) + cfg_dump = cfg.model_dump(exclude_none=True) + seq = _read_fasta_sequence(p) + ann = scorer.annotate(seq) + ann_view = [ + { + "type": getattr(a, "type", None), + "id": getattr(a, "id", None), + "start": getattr(a, "start", None), + "end": getattr(a, "end", None), + "strand": getattr(a, "strand", None), + } + for a in ann if getattr(a, "type", None) != "restriction_site" + ] + print( + f"testlog file={os.path.basename(p)} config={cfg_dump} score={score:.4f} components={components} annotations={ann_view}" + ) + + assert 0.0 <= score <= 1.0 + for key in ["ori", "promoter", "terminator", "marker", "cds", "length_factor"]: + assert key in components + assert 0.0 <= float(components[key]) <= 1.0 or key == "length_factor" + + +def test_location_bonus_non_decreasing(): + data_dir = os.path.join(os.path.dirname(__file__), "data") + path = os.path.join(data_dir, "pUC19.fasta") + seq = _read_fasta_sequence(path) + cfg_off = RewardConfig(location_aware=False) + cfg_on = RewardConfig(location_aware=True) + scorer_off = Scorer(cfg_off) + scorer_on = Scorer(cfg_on) + + score_off, comp_off = scorer_off.score(seq, source=os.path.basename(path)) + score_on, comp_on = scorer_on.score(seq, source=os.path.basename(path)) + ann = scorer_on.annotate(seq) + ann_view = [ + { + "type": getattr(a, "type", None), + "id": getattr(a, "id", None), + "start": getattr(a, "start", None), + "end": getattr(a, "end", None), + "strand": getattr(a, "strand", None), + } + for a in ann if getattr(a, "type", None) != "restriction_site" + ] + print(f"testlog file={os.path.basename(path)} config_off={cfg_off.model_dump()} score_off={score_off:.4f} comp_off={comp_off}") + print(f"testlog file={os.path.basename(path)} config_on={cfg_on.model_dump()} score_on={score_on:.4f} comp_on={comp_on} annotations={ann_view}") + + assert comp_on["cds"] >= comp_off["cds"] + assert score_on >= score_off + + +def test_wrapper_score_sequence_bounds_and_consistency(): + data_dir = os.path.join(os.path.dirname(__file__), "data") + path = os.path.join(data_dir, "pSC101.fasta") + seq = _read_fasta_sequence(path) + # wrapper returns 0..100 + s100 = rewards_mod.score_sequence(seq) + print(f"testlog file={os.path.basename(path)} wrapper_score_0_100={s100:.2f}") + assert 0.0 <= s100 <= 100.0 + + +def test_score_completions_handles_empty(): + data_dir = os.path.join(os.path.dirname(__file__), "data") + path = os.path.join(data_dir, "RF0G-IodoY.fasta") + seq = _read_fasta_sequence(path) + scores = rewards_mod.score_completions([seq, ""]) + print(f"testlog file={os.path.basename(path)} completions_scores={scores}") + assert len(scores) == 2 + assert scores[1] == 0.0 + assert 0.0 <= scores[0] <= 100.0 diff --git a/uv.lock b/uv.lock index 3a150df..a6f480c 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.11" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'linux'", @@ -1311,6 +1311,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/20/b0/36bd937216ec521246249be3bf9855081de4c5e06a0c9b4219dbeda50373/importlib_metadata-8.7.0-py3-none-any.whl", hash = "sha256:e5dd1551894c77868a30651cef00984d50e1002d06942a7101d34870c5f02afd", size = 27656, upload-time = "2025-04-27T15:29:00.214Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "interegular" version = "0.3.3" @@ -2880,7 +2889,7 @@ wheels = [ [[package]] name = "plasmidkit" version = "0.1.0" -source = { git = "https://github.com/McClain-Thiel/plasmid-kit.git#5342bc09d49c2c7377dfd436a3e2746bd1383a01" } +source = { git = "https://github.com/McClain-Thiel/plasmid-kit.git#3b799bd8e4c38115fe9df6147b5312df8403258d" } dependencies = [ { name = "biopython" }, { name = "edlib" }, @@ -2910,6 +2919,7 @@ dependencies = [ { name = "psutil" }, { name = "pydantic" }, { name = "pydantic-settings" }, + { name = "pytest" }, { name = "pyzmq" }, { name = "tensordict" }, { name = "torchrl", extra = ["llm"] }, @@ -2938,6 +2948,7 @@ requires-dist = [ { name = "psutil", specifier = ">=7.0.0" }, { name = "pydantic", specifier = ">=2.0.0" }, { name = "pydantic-settings", specifier = ">=2.0.0" }, + { name = "pytest", specifier = ">=8.4.2" }, { name = "pyzmq", specifier = ">=27.0.2" }, { name = "tensordict", specifier = ">=0.9.1" }, { name = "torchrl", extras = ["llm"], specifier = ">=0.9.2" }, @@ -2977,6 +2988,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/21/98/5ca173c8ec906abde26c28e1ecb34887343fd71cc4136261b90036841323/playwright-1.55.0-py3-none-win_arm64.whl", hash = "sha256:012dc89ccdcbd774cdde8aeee14c08e0dd52ddb9135bf10e9db040527386bd76", size = 31225543, upload-time = "2025-08-28T15:46:41.613Z" }, ] +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + [[package]] name = "prometheus-client" version = "0.23.1" @@ -3549,6 +3569,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/21/a0/1acdb8e2e2ad64e10ada85ad505126bfa3098ffb5b0f41a5bdcd9ef33b7b/pyrodigal-3.6.3.post1-cp313-cp313-win_amd64.whl", hash = "sha256:096dbdaae148e1c1a60b3d5bca84b7b5a7deb2d6a44b965dcc2cc640be02f667", size = 3307413, upload-time = "2025-03-04T00:54:54.077Z" }, ] +[[package]] +name = "pytest" +version = "8.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a3/5c/00a0e072241553e1a7496d638deababa67c5058571567b92a7eaa258397c/pytest-8.4.2.tar.gz", hash = "sha256:86c0d0b93306b961d58d62a4db4879f27fe25513d4b969df351abdddb3c30e01", size = 1519618, upload-time = "2025-09-04T14:34:22.711Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/a4/20da314d277121d6534b3a980b29035dcd51e6744bd79075a6ce8fa4eb8d/pytest-8.4.2-py3-none-any.whl", hash = "sha256:872f880de3fc3a5bdc88a11b39c9710c3497a547cfa9320bc3c5e62fbf272e79", size = 365750, upload-time = "2025-09-04T14:34:20.226Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" From 4a845a7bf1baa73a9335e585c10318f25e3d36b9 Mon Sep 17 00:00:00 2001 From: McClain Thiel Date: Fri, 31 Oct 2025 12:31:14 +0000 Subject: [PATCH 04/11] reorg and way better logging --- src/rewards/bioinformatics/__init__.py | 0 src/rewards/bioinformatics/logger.py | 73 +++++++++++++++++++ .../{ => bioinformatics}/reward_config.py | 0 src/rewards/{ => bioinformatics}/scorer.py | 9 ++- src/runners/grpo.py | 65 +++++++++++++++-- 5 files changed, 137 insertions(+), 10 deletions(-) create mode 100644 src/rewards/bioinformatics/__init__.py create mode 100644 src/rewards/bioinformatics/logger.py rename src/rewards/{ => bioinformatics}/reward_config.py (100%) rename src/rewards/{ => bioinformatics}/scorer.py (96%) diff --git a/src/rewards/bioinformatics/__init__.py b/src/rewards/bioinformatics/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/rewards/bioinformatics/logger.py b/src/rewards/bioinformatics/logger.py new file mode 100644 index 0000000..efa1904 --- /dev/null +++ b/src/rewards/bioinformatics/logger.py @@ -0,0 +1,73 @@ +from typing import Dict, List, Any +import numpy as np +from transformers import TrainerCallback +import wandb + + +class RewardComponentLogger(TrainerCallback): + """Log component-level rewards to W&B with buffering.""" + + def __init__(self, log_frequency: int = 10): + self.log_frequency = int(log_frequency) + self.component_buffer: Dict[str, List[float]] = { + "ori": [], + "promoter": [], + "terminator": [], + "marker": [], + "cds": [], + "length_factor": [], + "total_reward": [], + } + + def add_components(self, components: Dict[str, float], total_reward: float) -> None: + self.component_buffer["ori"].append(float(components["ori"])) + self.component_buffer["promoter"].append(float(components["promoter"])) + self.component_buffer["terminator"].append(float(components["terminator"])) + self.component_buffer["marker"].append(float(components["marker"])) + self.component_buffer["cds"].append(float(components["cds"])) + self.component_buffer["length_factor"].append(float(components["length_factor"])) + self.component_buffer["total_reward"].append(float(total_reward)) + + def on_step_end(self, args, state, control, **kwargs): + if self.log_frequency <= 0 or state.global_step == 0: + return + if (state.global_step % self.log_frequency) != 0: + return + if not self.component_buffer["total_reward"]: + return + + log_dict: Dict[str, Any] = {} + try: + for name, values in self.component_buffer.items(): + if not values: + continue + arr = np.asarray(values, dtype=np.float32) + base = f"reward_components/{name}" + log_dict[f"{base}/mean"] = float(np.mean(arr)) + log_dict[f"{base}/std"] = float(np.std(arr)) + log_dict[f"{base}/min"] = float(np.min(arr)) + log_dict[f"{base}/max"] = float(np.max(arr)) + + # Simple violation counters + ori_vals = self.component_buffer.get("ori", []) + marker_vals = self.component_buffer.get("marker", []) + cds_vals = self.component_buffer.get("cds", []) + if ori_vals: + log_dict["reward_violations/ori_perfect"] = float(sum(x >= 1.0 for x in ori_vals) / len(ori_vals)) + if marker_vals: + log_dict["reward_violations/marker_perfect"] = float(sum(x >= 1.0 for x in marker_vals) / len(marker_vals)) + if cds_vals: + log_dict["reward_violations/cds_zero"] = float(sum(x == 0.0 for x in cds_vals) / len(cds_vals)) + + # Periodic histograms + if (state.global_step % (self.log_frequency * 10)) == 0: + log_dict["reward_histograms/total_reward"] = wandb.Histogram(self.component_buffer["total_reward"]) # type: ignore[arg-type] + log_dict["reward_histograms/ori"] = wandb.Histogram(self.component_buffer["ori"]) # type: ignore[arg-type] + log_dict["reward_histograms/cds"] = wandb.Histogram(self.component_buffer["cds"]) # type: ignore[arg-type] + + wandb.log(log_dict, step=int(state.global_step)) + finally: + for key in self.component_buffer: + self.component_buffer[key] = [] + + diff --git a/src/rewards/reward_config.py b/src/rewards/bioinformatics/reward_config.py similarity index 100% rename from src/rewards/reward_config.py rename to src/rewards/bioinformatics/reward_config.py diff --git a/src/rewards/scorer.py b/src/rewards/bioinformatics/scorer.py similarity index 96% rename from src/rewards/scorer.py rename to src/rewards/bioinformatics/scorer.py index 8aab34f..2f5f71c 100644 --- a/src/rewards/scorer.py +++ b/src/rewards/bioinformatics/scorer.py @@ -68,7 +68,7 @@ def _to_feat(self, x: Any) -> "Scorer._Feat": evidence=getattr(x, "evidence", {}), ) - def _merge_group(self, feats: List[Any], threshold: float) -> List["Scorer._Feat"]: + def _merge_group(self, feats: List[Any], threshold: float, *, respect_strand: bool) -> List["Scorer._Feat"]: if not feats: return [] items = [self._to_feat(f) for f in feats] @@ -78,7 +78,8 @@ def _merge_group(self, feats: List[Any], threshold: float) -> List["Scorer._Feat for nxt in items[1:]: ovl = self._overlap_len(cur, nxt) min_len = max(1, min(self._len(cur), self._len(nxt))) - if ovl / float(min_len) >= threshold and cur.strand == nxt.strand: + strands_compatible = (cur.strand == nxt.strand) or (not respect_strand) + if ovl / float(min_len) >= threshold and strands_compatible: cur.start = min(cur.start, nxt.start) cur.end = max(cur.end, nxt.end) cur.id = f"{cur.id}|{nxt.id}" if cur.id or nxt.id else None @@ -102,7 +103,9 @@ def _preprocess_annotations(self, annotations: Any) -> List[Any]: merged_groups: Dict[str, List[Scorer._Feat]] = {} for t in ("rep_origin", "ori", "origin_of_replication", "promoter", "terminator", "marker", "cds"): if t in groups: - merged_groups[t] = self._merge_group(groups[t], thr) + # Ignore strand for ORI and marker; respect for others + respect = t not in ("rep_origin", "ori", "origin_of_replication", "marker") + merged_groups[t] = self._merge_group(groups[t], thr, respect_strand=respect) # suppress CDS if overlaps any non-CDS (ori/promoter/terminator/marker) non_cds: List[Scorer._Feat] = [] diff --git a/src/runners/grpo.py b/src/runners/grpo.py index 197dd31..54b735f 100644 --- a/src/runners/grpo.py +++ b/src/runners/grpo.py @@ -1,14 +1,19 @@ from datasets import load_dataset -from transformers import AutoTokenizer, set_seed +from transformers import AutoTokenizer, set_seed, TrainerCallback from trl import GRPOTrainer, GRPOConfig import torch from src.config import Config -from src.rewards import Score +from src.rewards.bioinformatics.scorer import Scorer +from src.rewards.bioinformatics.reward_config import RewardConfig +from src.rewards.bioinformatics.logger import RewardComponentLogger import datetime +from typing import List +import wandb -MODEL_ID = "McClain/plasmidgpt-addgene-gpt2" -TRAIN_PARQUET = "data/train.parquet" -VAL_PARQUET = "data/test.parquet" +cfg = Config() +MODEL_ID = cfg.model +TRAIN_PARQUET = cfg.train_dataset +VAL_PARQUET = cfg.val_dataset run_name = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") run_name = f"grpo-{run_name}" @@ -75,6 +80,7 @@ def select_prompt_and_extras(ds): save_steps=100, logging_strategy="steps", logging_steps=1, + report_to=["wandb"], bf16=torch.cuda.is_available(), gradient_checkpointing=False, max_grad_norm=0.5, @@ -92,7 +98,7 @@ def select_prompt_and_extras(ds): num_generations=8, max_completion_length=256, #min_completion_length=16, #not supported currently - temperature=0.80, + temperature=0.95, top_p=0.90, # GRPO specifics @@ -108,14 +114,59 @@ def select_prompt_and_extras(ds): vllm_mode="colocate" ) +# ---- reward config and scorer ---- +reward_config = RewardConfig( + ori_min=1, + ori_max=1, + promoter_min=1, + promoter_max=5, + terminator_min=0, + terminator_max=2, + marker_min=1, + marker_max=2, +) + +scorer = Scorer(reward_config) + +reward_logger = RewardComponentLogger(log_frequency=10) + + +def batch_reward_fn(samples: List[str], **kwargs) -> List[float]: + rewards: List[float] = [] + for seq in samples: + score, components = scorer.score(seq) + rewards.append(float(score)) + reward_logger.add_components(components, float(score)) + return rewards + + +# ---- W&B init ---- +wandb.init( + project=cfg.wandb_project, + entity=cfg.wandb_entity, + name=run_name, + config={ + "model": MODEL_ID, + "reward_config": reward_config.model_dump(), + "grpo_config": { + "beta": args.beta, + "epsilon": args.epsilon, + "temperature": args.temperature, + "num_generations": args.num_generations, + "loss_type": args.loss_type, + }, + }, +) + # ---- trainer (NO ref_model kwarg) ---- trainer = GRPOTrainer( model=MODEL_ID, - reward_funcs=Score, + reward_funcs=[batch_reward_fn], args=args, train_dataset=train_ds, eval_dataset=eval_ds if use_eval else None, processing_class=tok, + callbacks=[reward_logger], ) trainer.train() From 664a557c6e3ab060adf409e934669e77ede6a2ab Mon Sep 17 00:00:00 2001 From: McClain Thiel Date: Fri, 31 Oct 2025 12:39:09 +0000 Subject: [PATCH 05/11] added test --- .DS_Store | Bin 8196 -> 8196 bytes .github/workflows/tests.yml | 33 ++ README.md | 2 + config/.gitkeep | 0 config/naive_grpo.yaml | 152 --------- config/verl_grpo.yaml | 156 --------- config/verl_naive_ppo.yaml | 147 -------- config/verl_ppo.yaml | 154 --------- src/rewards/__init__.py | 9 +- src/rewards/bioinformatics/scorer.py | 2 +- src/rewards/plasmid_informatics.py | 479 --------------------------- src/rewards/rewards.py | 77 ----- src/rewards/verl_reward.py | 342 ------------------- tests/test_rewards.py | 30 +- 14 files changed, 55 insertions(+), 1528 deletions(-) create mode 100644 .github/workflows/tests.yml delete mode 100644 config/.gitkeep delete mode 100644 config/naive_grpo.yaml delete mode 100644 config/verl_grpo.yaml delete mode 100644 config/verl_naive_ppo.yaml delete mode 100644 config/verl_ppo.yaml delete mode 100644 src/rewards/plasmid_informatics.py delete mode 100644 src/rewards/rewards.py delete mode 100644 src/rewards/verl_reward.py diff --git a/.DS_Store b/.DS_Store index 1f447099abcc9847609503430c5335e257df3fec..406e62abdcb50d62a048548b0d33b108de0f3665 100644 GIT binary patch delta 33 pcmZp1XmOa}¥U^hRb!ekzS-<$gd<(N0C@;qYR%r5bl9RR>;3o-xz delta 162 zcmZp1XmOa}C7U^hRb%48ma-}P(^$qe}nc?@X`nGESU>4w3{`MCuQU|<~qq)_B@ v^Icq$a`Kaa;vA1f|2ZA?JMM_7CWWAyf()2FtRR~LHZuzhWg63z_(XI3dK diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..454ff6f --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,33 @@ +name: tests + +on: + push: + branches: [ main, master, develop ] + pull_request: + +jobs: + pytest: + runs-on: ubuntu-latest + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install uv + run: | + pip install --upgrade pip + pip install uv + + - name: Sync dependencies + run: | + uv sync + + - name: Run tests + run: | + uv run pytest -q + diff --git a/README.md b/README.md index b19e12d..082475a 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,7 @@ # PlasmidRL +[![Tests](https://github.com/McClain-Thiel/PlasmidRL/actions/workflows/tests.yml/badge.svg)](https://github.com/McClain-Thiel/PlasmidRL/actions/workflows/tests.yml) + Reinforcement learning experiments for plasmid design. Train language models to generate functional DNA sequences using various RL algorithms. ## Quick Start diff --git a/config/.gitkeep b/config/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/config/naive_grpo.yaml b/config/naive_grpo.yaml deleted file mode 100644 index 9e8e486..0000000 --- a/config/naive_grpo.yaml +++ /dev/null @@ -1,152 +0,0 @@ -# configs/grpo_minimal.yaml -defaults: - - ppo_trainer # using PPO trainer harness; GRPO is selected via algorithm.adv_estimator - - _self_ - -data: - tokenizer: McClain/naive-dna-llama-6mer - train_files: data/verl/train.parquet - val_files: data/verl/test.parquet - prompt_key: prompt - max_prompt_length: 1024 - max_response_length: 256 - train_batch_size: 4 - trust_remote_code: true - return_raw_input_ids: true - return_raw_chat: true - return_full_prompt: false - apply_chat_template_kwargs: - chat_template: | - {%- for m in messages -%} - {{- m['content'] -}} - {%- if not loop.last %}\n{%- endif -%} - {%- endfor -%} - shuffle: true - filter_overlong_prompts: false - filter_overlong_prompts_workers: 1 - truncation: right - dataloader_num_workers: 2 - custom_cls: - path: null - name: null - -actor_rollout_ref: - hybrid_engine: true - model: - path: McClain/naive-dna-llama-6mer - trust_remote_code: true - enable_gradient_checkpointing: false - override_config: - attn_implementation: eager - torch_dtype: bfloat16 - - actor: - strategy: fsdp - ppo_mini_batch_size: 4 - ppo_micro_batch_size_per_gpu: 1 - use_dynamic_bsz: false - ppo_max_token_len_per_gpu: 4096 - grad_clip: 1.0 - clip_ratio: 0.2 - entropy_coeff: 0.0 - use_kl_loss: false # KL via reward shaping for GRPO - tis_imp_ratio_cap: -1 - use_torch_compile: true - kl_loss_coef: 0.0 - kl_loss_type: low_var_kl - ppo_epochs: 1 - optim: - lr: 1.0e-5 - lr_warmup_steps_ratio: 0.0 - warmup_style: constant - total_training_steps: -1 - fsdp_config: - wrap_policy: - min_num_params: 0 - param_offload: false - optimizer_offload: false - model_dtype: bfloat16 - - ref: - log_prob_micro_batch_size_per_gpu: 4 - - rollout: - name: vllm - do_sample: true - top_p: 0.95 - temperature: 0.9 - n: 32 # multi-sample per prompt for GRPO - tensor_model_parallel_size: 1 - enforce_eager: true - free_cache_engine: false - gpu_memory_utilization: 0.6 - max_num_batched_tokens: 256 - max_num_seqs: 32 - enable_chunked_prefill: false - response_length: ${data.max_response_length} - log_prob_micro_batch_size_per_gpu: 4 - calculate_log_probs: false - val_kwargs: - do_sample: false - temperature: 0.0 - n: 1 - -critic: - enable: false # GRPO doesn't use a critic - model: - path: McClain/naive-dna-llama-6mer - trust_remote_code: true - fsdp_config: - wrap_policy: - min_num_params: 0 - param_offload: false - override_config: - attn_implementation: eager - optim: - lr: 5.0e-5 - ppo_micro_batch_size_per_gpu: 1 - -reward_model: - enable: false - -custom_reward_function: - path: src/rewards/verl_reward.py - name: get_plasmid_reward - -algorithm: - gamma: 1.0 - lam: 1.0 - adv_estimator: grpo # GRPO selected here - use_kl_in_reward: true # KL shaping instead of actor loss - kl_penalty: kl - kl_ctrl: - type: fixed - kl_coef: 0.001 - horizon: 10000 - target_kl: 0.1 - -trainer: - total_epochs: 1000 - project_name: verl-grpo - experiment_name: minimal - logger: ['console', 'wandb'] - log_val_generations: 0 - nnodes: 1 - n_gpus_per_node: 1 - save_freq: 50 - val_before_train: false - test_freq: 50 - critic_warmup: 0 - default_hdfs_dir: null - default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name} - resume_mode: auto - resume_from_path: null - remove_previous_ckpt_in_save: false - del_local_ckpt_after_load: false - ray_wait_register_center_timeout: 300 - wandb: - entity: mcclain - project: verl-grpo - name: naive-dna-grpo - tags: ['naive-dna', 'verl', 'grpo', 'rl'] - notes: 'GRPO training on naive-dna-llama-6mer for plasmid design' diff --git a/config/verl_grpo.yaml b/config/verl_grpo.yaml deleted file mode 100644 index ac001dd..0000000 --- a/config/verl_grpo.yaml +++ /dev/null @@ -1,156 +0,0 @@ -# configs/grpo_minimal.yaml -defaults: - - ppo_trainer - - _self_ - -data: - tokenizer: null - train_files: data/train.parquet - val_files: data/test.parquet - prompt_key: prompt - max_prompt_length: 1024 - max_response_length: 256 # shorter early; bump to 256 once stable - train_batch_size: 2 - trust_remote_code: true - return_raw_input_ids: true - return_raw_chat: true - return_full_prompt: false - apply_chat_template_kwargs: - chat_template: | - {%- for m in messages -%} - {{- m['content'] -}} - {%- if not loop.last %}\n{%- endif -%} - {%- endfor -%} - shuffle: true - filter_overlong_prompts: false - filter_overlong_prompts_workers: 1 - truncation: right - dataloader_num_workers: 2 - custom_cls: - path: null - name: null - -actor_rollout_ref: - hybrid_engine: true - - model: - path: McClain/plasmidgpt-addgene-gpt2 - trust_remote_code: true - enable_gradient_checkpointing: false - override_config: - attn_implementation: eager - _attn_implementation: eager - - actor: - strategy: fsdp - ppo_mini_batch_size: 2 - ppo_micro_batch_size_per_gpu: 1 - use_dynamic_bsz: false - ppo_max_token_len_per_gpu: 2048 - grad_clip: 0.5 - clip_ratio: 0.2 - entropy_coeff: 0.0 - - # GRPO: KL as a LOSS term (not in reward) - use_kl_loss: true - kl_loss_coef: 0.001 # tune 5e-4 – 2e-3 - kl_loss_type: low_var_kl - loss_agg_mode: token-mean - - use_torch_compile: false - ppo_epochs: 1 - - optim: - lr: 3e-6 # a bit more conservative to start - lr_warmup_steps_ratio: 0.0 - warmup_style: constant - total_training_steps: -1 - - fsdp_config: - wrap_policy: - min_num_params: 0 - param_offload: false - optimizer_offload: false - - # Explicit, frozen reference model for KL - ref: - model: - path: McClain/plasmidgpt-addgene-gpt2 - trust_remote_code: true - log_prob_micro_batch_size_per_gpu: 4 - - rollout: - name: vllm - do_sample: true - top_p: 0.90 - temperature: 0.80 - n: 9 - tensor_model_parallel_size: 1 - enforce_eager: true - free_cache_engine: false - gpu_memory_utilization: 0.15 - max_num_batched_tokens: 128 - max_num_seqs: 16 - enable_chunked_prefill: false - response_length: ${data.max_response_length} - log_prob_micro_batch_size_per_gpu: 4 - calculate_log_probs: false # not needed; loss path recomputes - val_kwargs: - do_sample: false - temperature: 0.0 - n: 1 - -critic: - enable: false - model: - path: McClain/plasmidgpt-addgene-gpt2 - trust_remote_code: true - fsdp_config: - wrap_policy: - min_num_params: 0 - param_offload: false - override_config: - attn_implementation: eager - optim: - lr: 2.5e-5 - ppo_micro_batch_size_per_gpu: 1 - -reward_model: - enable: false - -custom_reward_function: - path: src/rewards/verl_reward.py - name: compute_score - -algorithm: - gamma: 1.0 - lam: 1.0 - adv_estimator: grpo - use_kl_in_reward: false # << GRPO uses KL loss, not in-reward - # (remove kl_penalty / kl_ctrl entirely) - -trainer: - total_epochs: 20 - project_name: verl-grpo - experiment_name: grpo-with-markers-2 - logger: ['console', 'wandb'] - log_val_generations: 0 - nnodes: 1 - n_gpus_per_node: 1 - save_freq: 10 - val_before_train: false - test_freq: 10 - critic_warmup: 0 - default_hdfs_dir: null - default_local_dir: /s3/checkpoints/${trainer.project_name}/${trainer.experiment_name} - resume_mode: disable - resume_from_path: null - remove_previous_ckpt_in_save: false - del_local_ckpt_after_load: false - ray_wait_register_center_timeout: 300 - wandb: - entity: mcclain - project: verl-grpo - name: plasmidgpt-dna-grpo-better-loggings - tags: ['plasmidgpt-dna', 'verl', 'grpo', 'rl', 'long-run'] - notes: 'Extended GRPO training on plasmidgpt-dna for plasmid design' diff --git a/config/verl_naive_ppo.yaml b/config/verl_naive_ppo.yaml deleted file mode 100644 index c2e4b7c..0000000 --- a/config/verl_naive_ppo.yaml +++ /dev/null @@ -1,147 +0,0 @@ -# VERL PPO config for naive DNA LLaMA -# See architecture: https://verl.readthedocs.io/en/latest/examples/ppo_code_architecture.html - -defaults: - - ppo_trainer - - _self_ - -data: - tokenizer: McClain/naive-dna-llama-6mer - train_files: data/verl/train.parquet - val_files: data/verl/test.parquet - prompt_key: prompt - max_prompt_length: 512 - max_response_length: 1024 - train_batch_size: 4 - trust_remote_code: true - return_raw_input_ids: true - return_raw_chat: true - return_full_prompt: false - apply_chat_template_kwargs: - chat_template: | - {%- for m in messages -%} - {{- m['content'] -}} - {%- if not loop.last %}\n{%- endif -%} - {%- endfor -%} - shuffle: true - filter_overlong_prompts: false - filter_overlong_prompts_workers: 1 - truncation: right - dataloader_num_workers: 2 - custom_cls: - path: null - name: null - -actor_rollout_ref: - hybrid_engine: true - model: - path: McClain/naive-dna-llama-6mer - trust_remote_code: true - enable_gradient_checkpointing: false - override_config: {} - actor: - strategy: fsdp - ppo_mini_batch_size: 4 - ppo_micro_batch_size_per_gpu: 1 - use_dynamic_bsz: false - ppo_max_token_len_per_gpu: 4096 - grad_clip: 1.0 - clip_ratio: 0.2 - entropy_coeff: 0.0 - use_kl_loss: true - tis_imp_ratio_cap: -1 - use_torch_compile: true - kl_loss_coef: 0.001 - kl_loss_type: low_var_kl - ppo_epochs: 1 - optim: - lr: 1.0e-4 - lr_warmup_steps_ratio: 0.0 - warmup_style: constant - total_training_steps: -1 - fsdp_config: - wrap_policy: - min_num_params: 0 - param_offload: false - optimizer_offload: false - ref: - log_prob_micro_batch_size_per_gpu: 4 - rollout: - name: vllm - do_sample: true - top_p: 0.95 - temperature: 0.9 - n: 32 - tensor_model_parallel_size: 1 - enforce_eager: true - free_cache_engine: false - gpu_memory_utilization: 0.6 - max_num_batched_tokens: 256 - max_num_seqs: 32 - enable_chunked_prefill: false - response_length: ${data.max_response_length} - log_prob_micro_batch_size_per_gpu: 4 - calculate_log_probs: false - val_kwargs: - do_sample: false - temperature: 0.0 - n: 1 - -critic: - model: - path: McClain/naive-dna-llama-6mer - trust_remote_code: true - fsdp_config: - wrap_policy: - min_num_params: 0 - param_offload: false - override_config: {} - optim: - lr: 5.0e-5 - ppo_micro_batch_size_per_gpu: 1 - enable: false - -reward_model: - enable: false - -custom_reward_function: - path: src/rewards/verl_reward.py - name: get_plasmid_reward - -algorithm: - gamma: 1.0 - lam: 1.0 - adv_estimator: grpo - use_kl_in_reward: false - kl_penalty: kl - kl_ctrl: - type: fixed - kl_coef: 0.001 - horizon: 10000 - target_kl: 0.1 - -trainer: - total_epochs: 10 - project_name: verl-ppo - experiment_name: naive-llama - logger: ['console', 'wandb'] - log_val_generations: 0 - nnodes: 1 - n_gpus_per_node: 1 - save_freq: 50 - val_before_train: false - test_freq: 50 - critic_warmup: 0 - default_hdfs_dir: null - default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name} - resume_mode: auto - resume_from_path: null - remove_previous_ckpt_in_save: false - del_local_ckpt_after_load: false - ray_wait_register_center_timeout: 300 - wandb: - entity: mcclain - project: verl-ppo - name: plasmid-verl-naive-ppo - tags: ['plasmid', 'verl', 'ppo', 'naive-llama'] - notes: 'VERL PPO training for naive DNA LLaMA' diff --git a/config/verl_ppo.yaml b/config/verl_ppo.yaml deleted file mode 100644 index af2f630..0000000 --- a/config/verl_ppo.yaml +++ /dev/null @@ -1,154 +0,0 @@ -# configs/ppo_minimal.yaml -defaults: - - ppo_trainer - - _self_ -data: - tokenizer: null - train_files: data/verl/train.parquet - val_files: data/verl/test.parquet - prompt_key: prompt - max_prompt_length: 1024 - max_response_length: 256 - train_batch_size: 4 - trust_remote_code: true - return_raw_input_ids: true - return_raw_chat: true - return_full_prompt: false - apply_chat_template_kwargs: - chat_template: | - {%- for m in messages -%} - {{- m['content'] -}} - {%- if not loop.last %}\n{%- endif -%} - {%- endfor -%} - shuffle: true - filter_overlong_prompts: false - filter_overlong_prompts_workers: 1 - truncation: right - dataloader_num_workers: 2 - custom_cls: - path: null - name: null - -actor_rollout_ref: - hybrid_engine: true - model: - path: McClain/plasmidgpt-addgene-gpt2 - trust_remote_code: true - enable_gradient_checkpointing: false - override_config: - num_attention_heads: 12 - num_key_value_heads: 12 - attn_implementation: eager - # Optional PEFT: uncomment to prefer LoRA on small GPUs - # lora_rank: 16 - # lora_alpha: 32 - # lora_dropout: 0.05 - actor: - strategy: fsdp - ppo_mini_batch_size: 4 # match train_batch_size for single mini-batch - ppo_micro_batch_size_per_gpu: 1 - use_dynamic_bsz: false - ppo_max_token_len_per_gpu: 4096 - grad_clip: 1.0 - clip_ratio: 0.2 - entropy_coeff: 0.0 - use_kl_loss: true - tis_imp_ratio_cap: -1 - use_torch_compile: true - kl_loss_coef: 0.001 - kl_loss_type: low_var_kl - ppo_epochs: 1 - optim: - lr: 1.0e-5 - lr_warmup_steps_ratio: 0.0 - warmup_style: constant - total_training_steps: -1 - fsdp_config: - wrap_policy: - min_num_params: 0 - param_offload: false - optimizer_offload: false - ref: - log_prob_micro_batch_size_per_gpu: 4 - rollout: - name: vllm - do_sample: true - top_p: 0.95 - temperature: 0.9 - n: 32 - tensor_model_parallel_size: 1 - enforce_eager: true - free_cache_engine: false - gpu_memory_utilization: 0.6 - max_num_batched_tokens: 256 - max_num_seqs: 32 - enable_chunked_prefill: false - response_length: ${data.max_response_length} - log_prob_micro_batch_size_per_gpu: 4 - calculate_log_probs: false - val_kwargs: - do_sample: false - temperature: 0.0 - n: 1 - -critic: - model: - path: McClain/plasmidgpt-addgene-gpt2 - trust_remote_code: true - fsdp_config: - wrap_policy: - min_num_params: 0 - param_offload: false - override_config: - num_attention_heads: 12 - num_key_value_heads: 12 - attn_implementation: eager - optim: - lr: 5.0e-5 - ppo_micro_batch_size_per_gpu: 1 - enable: false - -reward_model: - enable: false # using your custom function instead - -custom_reward_function: - path: src/rewards/verl_reward.py - name: get_plasmid_reward - -algorithm: - gamma: 1.0 - lam: 1.0 - adv_estimator: grpo - use_kl_in_reward: false - kl_penalty: kl - kl_ctrl: - type: fixed - kl_coef: 0.001 - horizon: 10000 - target_kl: 0.1 - -trainer: - total_epochs: 10 # more epochs since we have limited data diversity - project_name: verl-ppo - experiment_name: minimal - logger: ['console', 'wandb'] - log_val_generations: 0 - nnodes: 1 - n_gpus_per_node: 1 - save_freq: 50 - val_before_train: false - test_freq: 50 - critic_warmup: 0 - default_hdfs_dir: null - default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name} - resume_mode: auto - resume_from_path: null - remove_previous_ckpt_in_save: false - del_local_ckpt_after_load: false - ray_wait_register_center_timeout: 300 - wandb: - entity: mcclain - project: verl-ppo - name: plasmid-verl-ppo - tags: ['plasmid', 'verl', 'ppo', 'rl'] - notes: 'VERL PPO training for plasmid design optimization' diff --git a/src/rewards/__init__.py b/src/rewards/__init__.py index f82eada..7d958f1 100644 --- a/src/rewards/__init__.py +++ b/src/rewards/__init__.py @@ -1,8 +1,5 @@ -from .scorer import Scorer -from .reward_config import RewardConfig - -config = RewardConfig( - -) +from .bioinformatics.scorer import Scorer +from .bioinformatics.reward_config import RewardConfig +# Convenience alias: call as `from src.rewards import score` if needed score = Scorer(RewardConfig()).score \ No newline at end of file diff --git a/src/rewards/bioinformatics/scorer.py b/src/rewards/bioinformatics/scorer.py index 2f5f71c..5b7b9b0 100644 --- a/src/rewards/bioinformatics/scorer.py +++ b/src/rewards/bioinformatics/scorer.py @@ -1,4 +1,4 @@ -from src.rewards.reward_config import RewardConfig +from src.rewards.bioinformatics.reward_config import RewardConfig import plasmidkit as pk from typing import Any, List, Tuple, Union, Dict from concurrent.futures import ThreadPoolExecutor diff --git a/src/rewards/plasmid_informatics.py b/src/rewards/plasmid_informatics.py deleted file mode 100644 index fc63a34..0000000 --- a/src/rewards/plasmid_informatics.py +++ /dev/null @@ -1,479 +0,0 @@ -import logging -import time -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Optional, List, Dict, Tuple - -import httpx -import requests -import torch -from tensordict import TensorDict -from tensordict.tensorclass import NonTensorStack, NonTensorData -from torchrl.data.tensor_specs import CompositeSpec, Unbounded -from torchrl.envs import Transform - -from ..config import config - -logger = logging.getLogger(__name__) - -class RewardTransform(Transform): - """Assign rewards by calling external scoring endpoints in parallel.""" - - def __init__( - self, - rewards_server_url: Optional[str] = None, - timeout_s: float = 60.0, - max_workers: Optional[int] = None, - log_timings: bool = False, - ): - super().__init__(in_keys=[], out_keys=["reward"]) - - self.rewards_server_url = (rewards_server_url or config.informatics_server_url).rstrip("/") - self.timeout_s = timeout_s - self.client = httpx.Client( - base_url=self.rewards_server_url, - timeout=httpx.Timeout(timeout_s), - follow_redirects=True, - ) - #defaults for evaluating plasmif - self.require = {"ori": None, "amr": None, "mcs": True, "promoter": None} - self.weights = {"ori": 0.30, "amr": 0.30, "mcs": 0.20, "promoter": 0.20} - self.gc = {"target": 0.55, "weight": 0.05, "tolerance": 0.10} - # length bonus configuration: up to 0.25 at or below target length, decays slightly for longer sequences - self.length = {"target": 1000, "weight": 0.25, "decay_per_bp": 1e-4, "penalize_shorter": False} - self._test_connection() - - # Exact paths & params per your curl commands - self._endpoints: List[Dict] = [ - {"name": "amrfinder", "path": "/amrfinder/text", "params": {"is_protein": "false", "format": "json"}}, - {"name": "prodigal", "path": "/prodigal/text", "params": {"mode": "auto", "format": "json"}}, - {"name": "plannotate","path": "/plannotate/fast","params": {}}, - ] - - # Plain text in, accept json or text back - self._headers = { - "Content-Type": "text/plain; charset=utf-8", - "Accept": "application/json, text/plain; q=0.9, */*; q=0.1", - } - self._max_workers = max_workers or max(1, len(self._endpoints)) - self._log_timings = log_timings - self._sample_executor = ThreadPoolExecutor(max_workers=self._max_workers) - - def _test_connection(self): - try: - r = requests.get(self.rewards_server_url + "/health", timeout=5) - r.raise_for_status() - except requests.RequestException as e: - raise RuntimeError(f"Failed to connect to rewards server: {e}") - - def _post_text(self, path: str, params: Dict, text: str, name: str) -> Dict: - try: - resp = self.client.post( - path, params=params, content=text.encode("utf-8"), headers=self._headers - ) - ok = (resp.status_code == 200) - if not ok: - snippet = text[:120] - if len(text) > 120: - snippet += "…" - try: - preview = (resp.text or "")[:300] - except Exception: - preview = "" - logger.error( - "[%s] %s -> %s | %s | payload_snippet=%s", - name, - path, - resp.status_code, - preview, - snippet, - ) - try: - response_data = resp.json() - except Exception as json_err: - if ok: - logger.warning("[%s] %s JSON parse error: %s", name, path, json_err) - response_data = {} - return {"status": ok, "name": name, "reponse": response_data} - except Exception as e: - logger.error("[%s] %s call failed: %s", name, path, e) - return {"status": False, "name": name, "reponse": {}} - - def _score_single(self, idx: int, text: str, overrides: dict | None) -> tuple[int, float]: - if not text: - return idx, 0.0 - - if logger.isEnabledFor(logging.DEBUG): - logger.debug( - "Scoring sample idx=%s length=%s overrides=%s", - idx, - len(text), - bool(overrides), - ) - - with ThreadPoolExecutor(max_workers=len(self._endpoints)) as executor: - future_to_cfg = { - executor.submit( - self._post_text, - cfg["path"], - cfg.get("params", {}), - text, - cfg["name"], - ): cfg - for cfg in self._endpoints - } - results = [] - for future in as_completed(future_to_cfg): - try: - result = future.result() - except Exception as e: - result = {"status": False, "name": "", "reponse": {}, "error": str(e)} - results.append(result) - - try: - reward = self.combine_rewards(results, overrides=overrides) - except Exception as e: - logger.warning("Reward calculation failed for text: %s", e) - reward = 0.0 - - # Apply length-based component (string length, not tokens) - # Maximum of 0.25 when sequence length is at or below target; gently decays for longer sequences. - try: - lconf = dict(self.length) - if overrides and isinstance(overrides, dict) and "length" in overrides: - lconf.update(overrides.get("length") or {}) - - target_len = max(0, int(lconf.get("target", 1000))) - weight = float(lconf.get("weight", 0.25)) - decay_per_bp = float(lconf.get("decay_per_bp", 1e-4)) - penalize_shorter = bool(lconf.get("penalize_shorter", False)) - - seq_len = len(text) - if seq_len <= target_len and not penalize_shorter: - length_bonus = weight - else: - delta = max(0, seq_len - target_len) if not penalize_shorter else abs(seq_len - target_len) - factor = max(0.0, 1.0 - decay_per_bp * float(delta)) - length_bonus = weight * factor - reward += float(length_bonus) - except Exception: - # If anything goes wrong with length bonus, skip it silently - pass - - if logger.isEnabledFor(logging.DEBUG): - status_map = { - str(res.get("name", "?")): bool(res.get("status", False)) - for res in results - } - logger.debug( - "Sample idx=%s component_status=%s reward=%.4f", - idx, - status_map, - reward, - ) - - return idx, reward - - def combine_rewards(self, info_dicts: List[Dict], overrides: dict | None = None) -> float: - """ - Combine amrfinder, prodigal, and plannotate responses into a single reward. - Scoring: - - Presence of ORI, AMR gene, MCS, promoter (weighted), each multiplied by percent identity if available. - - Small GC proximity bonus (around 55%) as a tie-breaker. - Extensible via self.require / self.weights / self.gc. - """ - require = dict(self.require) - weights = dict(self.weights) - gcconf = dict(self.gc) - - # Merge per-example overrides if provided - if overrides: - if "require" in overrides: require.update(overrides["require"] or {}) - if "weights" in overrides: weights.update(overrides["weights"] or {}) - if "gc" in overrides: gcconf.update(overrides["gc"] or {}) - - # filter out failed calls (False) and normalize shape - records = [r for r in info_dicts if isinstance(r, dict) and r.get("status") is not False] - by_name = {r.get("name"): (r.get("reponse") or {}) for r in records} - - amr = by_name.get("amrfinder") or {} - prodigal = by_name.get("prodigal") or {} - plannotate = by_name.get("plannotate") or [] - - # helpers - def _lc(s): return str(s or "").lower() - def _contains_any(text: str, needles) -> bool: - t = _lc(text) - return any(_lc(n) in t for n in (needles or [])) - def _best_pident(entries) -> float: - """Return best percent identity (0..1) among provided plannotate entries with 'pident'.""" - best = 0.0 - for e in entries: - try: - p = float(e.get("pident", 100.0)) / 100.0 - except Exception: - p = 1.0 - best = max(best, max(0.0, min(1.0, p))) - return best if best > 0 else 1.0 # default to 1.0 if none provided - - # ---- ORI / MCS / Promoter from plannotate ---- - ori_present, mcs_present, prom_present = False, False, False - ori_pident, mcs_pident, prom_pident = 1.0, 1.0, 1.0 - - if isinstance(plannotate, list): - # buckets - ori_entries, mcs_entries, prom_entries = [], [], [] - - for feat in plannotate: - name = str(feat.get("Feature") or "") - desc = str(feat.get("Description") or "") - typ = str(feat.get("Type") or "") - text = f"{name} {desc} {typ}" - - # ORI - if self.require["ori"] is None: - if typ.lower() == "rep_origin" or _contains_any(text, ["ori", "colE1", "pmb1", "pbr322", "puc"]): - ori_entries.append(feat) - else: - if _contains_any(text, self.require["ori"]): - ori_entries.append(feat) - - # MCS - if self.require["mcs"] is True: - if _contains_any(text, ["mcs", "multiple cloning site"]): - mcs_entries.append(feat) - elif isinstance(self.require["mcs"], list): - if _contains_any(text, self.require["mcs"]): - mcs_entries.append(feat) - - # Promoter - if self.require["promoter"] is None: - if typ.lower() == "promoter" or _contains_any(text, ["promoter"]): - prom_entries.append(feat) - else: - if _contains_any(text, self.require["promoter"]): - prom_entries.append(feat) - - if ori_entries: - ori_present = True - ori_pident = _best_pident(ori_entries) - if mcs_entries: - mcs_present = True - mcs_pident = _best_pident(mcs_entries) - if prom_entries: - prom_present = True - prom_pident = _best_pident(prom_entries) - - # ---- AMR from amrfinder ---- - amr_present, amr_pident = False, 1.0 - if isinstance(amr, dict): - hits = amr.get("genes", []) or [] - candidate_hits = [] - for g in hits: - cls = str(g.get("class") or "") - sym = str(g.get("element_symbol") or "") - nm = str(g.get("element_name") or "") - hay = f"{cls} {sym} {nm}" - if self.require["amr"] is None or _contains_any(hay, self.require["amr"]): - candidate_hits.append(g) - if candidate_hits: - amr_present = True - # pick best identity among matching hits - best = 1.0 - for g in candidate_hits: - try: - pid = float(g.get("percent_identity_to_reference", 100.0)) / 100.0 - except Exception: - pid = 1.0 - best = max(best, max(0.0, min(1.0, pid))) - amr_pident = best - - # ---- main score (presence × weight × identity) ---- - w = self.weights - main = 0.0 - if ori_present: main += w.get("ori", 0.0) * float(ori_pident) - if amr_present: main += w.get("amr", 0.0) * float(amr_pident) - if mcs_present: main += w.get("mcs", 0.0) * float(mcs_pident) - if prom_present: main += w.get("promoter", 0.0) * float(prom_pident) - main = min(main, 1.0) # keep things tidy - - # ---- GC tie-breaker from prodigal ---- - gc_bonus = 0.0 - if isinstance(prodigal, dict): - meta = prodigal.get("metadata", {}) or {} - gc_raw = meta.get("model_gc_cont") or meta.get("gc_cont") - if gc_raw is not None: - try: - s = str(gc_raw).strip().replace("%", "") - v = float(s) - gc = v / 100.0 if "%" in str(gc_raw) or v > 1.0 else v - target = float(self.gc["target"]) - tol = max(1e-6, float(self.gc["tolerance"])) - dist = abs(gc - target) - norm = max(0.0, 1.0 - (dist / tol)) # 1 at target, 0 at >= tolerance - gc_bonus = float(self.gc["weight"]) * norm - except Exception: - gc_bonus = 0.0 - - total = main + gc_bonus - return float(total) - - - def _call(self, td: TensorDict) -> TensorDict: - def _extract_strings(obj) -> list[str]: - if obj is None: - return [] - if isinstance(obj, str): - return [obj] - if isinstance(obj, NonTensorData): - return _extract_strings(obj.data) - if isinstance(obj, NonTensorStack): - values: list[str] = [] - for item in obj: - values.extend(_extract_strings(item)) - return values - if isinstance(obj, (list, tuple)): - values: list[str] = [] - for item in obj: - values.extend(_extract_strings(item)) - return values - return [str(obj)] - - llm_texts: list[str] = [] - keys = td.keys(True) - if ("text", "response") in keys: - llm_texts = _extract_strings(td.get(("text", "response"))) - elif ("text", "full") in keys: - llm_texts = _extract_strings(td.get(("text", "full"))) - elif "text" in keys: - llm_texts = _extract_strings(td.get("text")) - elif "query" in keys: - llm_texts = _extract_strings(td.get("query")) - - if not llm_texts: - zero = torch.zeros(td.batch_size + (1,), dtype=torch.float32) - if td.device is not None: - zero = zero.to(td.device) - td["reward"] = zero - logger.debug("RewardTransform received empty text; assigning zero reward") - return td - - # Align to batch size if we only got a single flattened string - expected = td.numel() - if len(llm_texts) != expected: - if len(llm_texts) == 1: - llm_texts = llm_texts * expected - else: - llm_texts = llm_texts[:expected] - - def _clean_sequence(raw: str) -> str: - text = str(raw) - # split FastChat-style or OpenAI-style special tokens - tokens = text.split("<|im_start|>") - segments = [] - for tok in tokens: - if not tok: - continue - if "|>" in tok: - _, remainder = tok.split("|>", 1) - else: - remainder = tok - segments.append(remainder) - text = "".join(segments) - # Remove explicit end tokens - text = text.replace("<|im_end|>", "") - # Keep only nucleotide characters and standard ambiguity codes - allowed = set("ACGTURYKMSWBDHVNXacgturykmswbdhvnx") - cleaned = "".join(ch for ch in text if ch in allowed) - cleaned = cleaned.upper() - # Ensure we don't return an empty sequence - return cleaned - - llm_texts = [_clean_sequence(s) for s in llm_texts] - - # 2) extract per-example overrides (align shape) - overrides = td.get("reward_params", None) - if overrides is None or isinstance(overrides, dict): - overrides_list = [overrides] * len(llm_texts) - else: - overrides_list = list(overrides) - if len(overrides_list) != len(llm_texts): - overrides_list = [None] * len(llm_texts) - - start_time = time.perf_counter() if self._log_timings else None - - futures = [] - for idx, (text, ov) in enumerate(zip(llm_texts, overrides_list)): - future = self._sample_executor.submit(self._score_single, idx, text, ov) - futures.append(future) - rewards: list[float] = [0.0] * len(futures) - failures = 0 - for future in futures: - try: - idx, reward_val = future.result() - rewards[idx] = float(reward_val) - except Exception as exc: - failures += 1 - logger.warning("Reward computation task failed: %s", exc) - - successes = len(rewards) - failures - if logger.isEnabledFor(logging.DEBUG): - preview = ", ".join(f"{val:.3f}" for val in rewards[:3]) - logger.debug( - "RewardTransform batch complete: size=%d successes=%d failures=%d preview=[%s]", - len(rewards), - successes, - failures, - preview, - ) - - if self._log_timings and start_time is not None: - elapsed = time.perf_counter() - start_time - logger.info( - "[RewardTransform] processed %d sequences in %.2fs (failures=%d)", - len(rewards), - elapsed, - failures, - ) - - # 4) write back (vector or scalar) - out = torch.as_tensor(rewards, dtype=torch.float32) - # reshape to td batch - if out.ndim == 1 and out.numel() == td.numel(): - out = out.view(td.batch_size + (1,)) - else: - out = out.mean().view(td.batch_size + (1,)) # conservative fallback - if td.device is not None: - out = out.to(td.device) - td["reward"] = out - return td - - - def transform_reward_spec(self, reward_spec: CompositeSpec) -> CompositeSpec: - device = getattr(reward_spec, "device", torch.device("cpu")) - reward_spec["reward"] = Unbounded( - shape=reward_spec.shape + (1,), - dtype=torch.float32, - device=device, - ) - return reward_spec - - def close(self): - try: - self.client.close() - except Exception: - pass - try: - self._sample_executor.shutdown(wait=False, cancel_futures=True) - except Exception: - pass - - def __del__(self): - try: - if not self.client.is_closed: - self.client.close() - except Exception: - pass - try: - self._sample_executor.shutdown(wait=False, cancel_futures=True) - except Exception: - pass diff --git a/src/rewards/rewards.py b/src/rewards/rewards.py deleted file mode 100644 index a60f2b8..0000000 --- a/src/rewards/rewards.py +++ /dev/null @@ -1,77 +0,0 @@ -from concurrent.futures import ThreadPoolExecutor -from functools import partial -from typing import Any, List -import logging -import time - -import plasmidkit as pk - -from src.rewards.scorer import Scorer -from src.rewards.reward_config import RewardConfig - - -logger = logging.getLogger("reward_logger") - -# Toggle detailed timing logs -try: - from src.config import Config - _REWARD_LOG_TIMINGS = bool(Config().reward_log_timings) -except Exception: - _REWARD_LOG_TIMINGS = False - -_SCORER = Scorer(RewardConfig()) - - -def annotate_completions(completions: list[str]) -> list[Any]: - """Annotate a flat list of completions using threads; strips spaces from sequences.""" - if _REWARD_LOG_TIMINGS: - t0 = time.perf_counter() - sequences = [s.replace(" ", "") for s in completions] - with ThreadPoolExecutor() as executor: - annotate = partial(pk.annotate, is_sequence=True) - annotations = list(executor.map(annotate, sequences)) - if _REWARD_LOG_TIMINGS: - dt_ms = (time.perf_counter() - t0) * 1000.0 - logger.info(f"reward.annotate n={len(completions)} time_ms={dt_ms:.2f}") - return annotations - - -def score_sequence(sequence: str) -> float: - """Config-driven score in [0, 100] using Scorer (auto-annotates).""" - if not sequence or not str(sequence).strip(): - return 0.0 - score01, _ = _SCORER.score(sequence) - return float(max(0.0, min(100.0, 100.0 * score01))) - - -def score_completions(completions: list[str]) -> list[float]: - if _REWARD_LOG_TIMINGS: - t0 = time.perf_counter() - if not completions: - logger.warning("reward.score_completions called with empty completions list") - return [] - - scores: List[float] = [] - non_empty_indices: List[int] = [] - non_empty_completions: List[str] = [] - - for i, c in enumerate(completions): - if not c or len(c.strip()) == 0: - scores.append(0.0) - else: - non_empty_indices.append(i) - non_empty_completions.append(c) - scores.append(0.0) # placeholder - - if non_empty_completions: - for idx in non_empty_indices: - scores[idx] = score_sequence(completions[idx]) - - if _REWARD_LOG_TIMINGS: - dt_ms = (time.perf_counter() - t0) * 1000.0 - n = len(scores) - mean_score = sum(scores) / n if n else 0.0 - logger.info( - f"reward.score n={n} mean={mean_score:.2f} time_ms={dt_ms:.2f}" - ) - return scores diff --git a/src/rewards/verl_reward.py b/src/rewards/verl_reward.py deleted file mode 100644 index b4da315..0000000 --- a/src/rewards/verl_reward.py +++ /dev/null @@ -1,342 +0,0 @@ -""" -VERL-compatible reward function using local plasmidkit annotations. -Adapted from src/rewards/rewards.py to work with VERL's expected interface. -""" - -import logging -import time -from concurrent.futures import ThreadPoolExecutor -from functools import partial -from typing import Any, Iterable, List, Tuple, Optional, Dict - -try: - import plasmidkit as pk -except ImportError: - pk = None - -try: - import wandb -except ImportError: - wandb = None - -logger = logging.getLogger(__name__) - -# Toggle detailed timing logs -try: - from src.config import Config - _REWARD_LOG_TIMINGS = bool(Config().reward_log_timings) -except Exception: - _REWARD_LOG_TIMINGS = False - - -def annotate_completions(completions: List[str]) -> List[Any]: - """Annotate a flat list of completions using threads; strips spaces from sequences.""" - if _REWARD_LOG_TIMINGS: - t0 = time.perf_counter() - - - - sequences = [s.replace(" ", "") for s in completions] - with ThreadPoolExecutor() as executor: - annotate = partial(pk.annotate, is_sequence=True) - annotations = list(executor.map(annotate, sequences)) - - if _REWARD_LOG_TIMINGS: - dt_ms = (time.perf_counter() - t0) * 1000.0 - logger.info(f"reward.annotate n={len(completions)} time_ms={dt_ms:.2f}") - - return annotations - - -def score_sequence( - sequence: str, - annotations: Any, - target_keywords: Iterable[str] | None = None, - return_components: bool = False, -) -> float | Tuple[float, Dict[str, float]]: - """ - Heuristic backbone score for GRPO. - - Rewards: - • ORI present: +20 for first ORI, -10 for each additional ORI. - • Marker present: +10 for having at least one marker. - • Up to two highest-scoring cassettes with partial credit: - - promoter→CDS: order (+5) + proximity (≤100bp:+5, ≤300:+3, ≤500:+2, ≤1000:+1). - - CDS→terminator: order (+5) + proximity (same as above). - - Out-of-order legs get proximity-only partial (+≤3). - - +2 if CDS is a marker and promoter is within 300 bp. - • Standalone promoters (+1 each up to +5). - • Standalone terminators (+1 each up to +5). - • Payload (GOI) CDS anywhere: +8, plus +4 if any promoter within 500 bp. - • Length bonus: shorter sequences favored (linear from +10 at ≤5kb to 0 at ≥30kb). - - Args: - sequence: DNA sequence to score - annotations: Plasmidkit annotations - target_keywords: Keywords for identifying payload genes - return_components: If True, return (score, components_dict) - - Returns: - float score [0, 100] if return_components=False - Tuple[float, Dict[str, float]] if return_components=True - """ - # ---- collect features (case-insensitive 'type') ---- - def T(x): return (x.type or "").lower() - feats = list(annotations) - oris = [x for x in feats if T(x) in ("rep_origin", "ori", "origin_of_replication")] - promoters = [x for x in feats if T(x) == "promoter"] - cdss = [x for x in feats if T(x) == "cds"] - terminators = [x for x in feats if T(x) == "terminator"] - markers = [x for x in feats if T(x) == "marker"] - - # ---- small helpers ---- - def strand(x): return getattr(x, "strand", "+") - def start(x): return int(getattr(x, "start", 0)) - def end(x): return int(getattr(x, "end", 0)) - def role(x): return (getattr(x, "evidence", {}) or {}).get("role", "").lower() - - def in_order_same_strand(a, b) -> bool: - if strand(a) != strand(b): return False - if strand(a) == "+": return start(a) <= start(b) - else: return end(a) >= end(b) - - def distance(a, b) -> int: - if end(a) < start(b): return start(b) - end(a) - if end(b) < start(a): return start(a) - end(b) - return 0 # overlapping/adjacent - - def prox_points(d: int, max_pts: int = 5) -> int: - if d <= 100: return max_pts - if d <= 300: return max_pts - 2 # 3 - if d <= 500: return max_pts - 3 # 2 - if d <= 1000: return 1 - return 0 - - # ---- ORI scoring ---- - score = 0.0 - components = {} - - ori_score = 0.0 - if len(oris) >= 1: - ori_score += 20.0 # +20 for first ORI - if len(oris) > 1: - ori_score -= 10.0 * (len(oris) - 1) # -10 for each additional ORI - score += ori_score - components["ori"] = ori_score - components["ori_count"] = float(len(oris)) - - # ---- Marker bonus ---- - marker_score = 0.0 - if len(markers) >= 1: - marker_score = 10.0 # +10 for having at least one marker - score += marker_score - components["marker"] = marker_score - components["marker_count"] = float(len(markers)) - - # ---- Cassette scoring (incl. out-of-order partials) ---- - def best_cassettes() -> List[Tuple[Any, Any, Any, int]]: - triples = [] - for p in promoters: - cds_same = [c for c in cdss if strand(c) == strand(p)] - if not cds_same: continue - cds_same.sort(key=lambda c: distance(p, c)) - c = cds_same[0] - - term_cands = [t for t in terminators if strand(t) == strand(c)] - # prefer ordered downstream terminators, but allow out-of-order partial credit - term_cands.sort(key=lambda t: distance(c, t)) - t = term_cands[0] if term_cands else None - - pts = 0 - # promoter -> CDS - if in_order_same_strand(p, c): - pts += 5 - pts += prox_points(distance(p, c), 5) - else: - pts += min(3, prox_points(distance(p, c), 5)) # out-of-order partial - - # CDS -> terminator - if t is not None: - if in_order_same_strand(c, t): - pts += 5 - pts += prox_points(distance(c, t), 5) - else: - pts += min(3, prox_points(distance(c, t), 5)) # out-of-order partial - - if role(c) == "marker" and distance(p, c) <= 300: - pts += 2 - - triples.append((p, c, t, pts)) - triples.sort(key=lambda x: x[3], reverse=True) - return triples[:2] - - cassette_score = 0.0 - for (_, _, _, pts) in best_cassettes(): - cassette_score += float(min(20, pts)) - score += cassette_score - components["cassette"] = cassette_score - - # ---- Standalone promoter & terminator credit ---- - promoter_score = 0.0 - if promoters: - promoter_score = float(min(5.0, 1.0 * len(promoters))) # +1 each up to +5 - score += promoter_score - components["promoter"] = promoter_score - components["promoter_count"] = float(len(promoters)) - - terminator_score = 0.0 - if terminators: - terminator_score = float(min(5.0, 1.0 * len(terminators))) # +1 each up to +5 - score += terminator_score - components["terminator"] = terminator_score - components["terminator_count"] = float(len(terminators)) - - # ---- Payload (GOI) anywhere ---- - target_keywords = [k.lower() for k in (target_keywords or [])] - def is_payload(c) -> bool: - r = role(c) - id_l = (getattr(c, "id", "") or "").lower() - return r in {"payload", "goi", "reporter"} or (target_keywords and any(k in id_l for k in target_keywords)) - - payloads = [c for c in cdss if is_payload(c)] - payload_score = 0.0 - if payloads: - payload_score = 8.0 - # +4 if any promoter within 500 bp (either direction; same strand preferred implicitly by proximity) - if any(distance(p, c) <= 500 for c in payloads for p in promoters): - payload_score += 4.0 - score += payload_score - components["payload"] = payload_score - components["payload_count"] = float(len(payloads)) - components["cds_count"] = float(len(cdss)) - - # ---- Length reward: shorter sequences get bonus ---- - L = max(0, len(sequence or "")) - - def length_reward(L: int) -> float: - # Linear reward from 10 pts at 5kb down to 0 pts at 30kb - # Sequences > 30kb get 0 reward - # Sequences <= 5kb get max 10 pts - if L >= 30000: - return 0.0 - if L <= 5000: - return 10.0 - # Linear interpolation between 5kb and 30kb - return 10.0 * (30000 - L) / (30000 - 5000) - - len_score = length_reward(L) - score += len_score - components["length_bonus"] = len_score - components["sequence_length"] = float(L) - - # ---- Clamp ---- - final_score = float(max(0.0, min(100.0, score))) - components["total_score"] = final_score - - if return_components: - return final_score, components - return final_score - - -def extract_dna_sequence(solution_str: str) -> str: - """ - Extract DNA sequence from the generated solution string. - Assumes the model generates DNA sequences (ACGT characters). - Strips whitespace and non-ACGT characters. - """ - if not solution_str: - return "" - - # Remove whitespace - seq = solution_str.strip().replace(" ", "").replace("\n", "").replace("\r", "") - - # Filter to only valid DNA bases (case-insensitive) - valid_bases = set("ACGTacgt") - seq = "".join(c for c in seq if c in valid_bases) - - return seq.upper() - - -def compute_score( - data_source: Optional[str], - solution_str: Optional[str], - ground_truth: Optional[str], - extra_info: Optional[Dict] = None, -) -> float: - """ - VERL-compatible reward function. - - This is the main entry point called by VERL's RewardManager. - - Args: - data_source: Dataset name (unused in this implementation) - solution_str: Generated text from the model (should contain DNA sequence) - ground_truth: Ground truth from the dataset (unused in this implementation) - extra_info: Additional info dict (unused in this implementation) - - Returns: - float: Reward score, normalized to [0, 1] range (from 0-100 internal score) - """ - if _REWARD_LOG_TIMINGS: - t0 = time.perf_counter() - - # Extract DNA sequence from solution - sequence = extract_dna_sequence(solution_str or "") - - if not sequence: - logger.warning("Empty or invalid DNA sequence in solution_str") - return 0.0 - - - annotations = pk.annotate(sequence, is_sequence=True) - - # Score the sequence (returns 0-100) with component breakdown - raw_score, components = score_sequence(sequence, annotations, return_components=True) - - # Normalize to [0, 1] for VERL - normalized_score = raw_score / 100.0 - - # Log reward components for monitoring (at INFO level so they appear in console) - # Sample logging: only log ~1% of the time to avoid flooding - import random - if random.random() < 0.01 or _REWARD_LOG_TIMINGS: - logger.info( - f"REWARD_BREAKDOWN: total={raw_score:.2f} | " - f"ori={components.get('ori', 0):.1f}(n={int(components.get('ori_count', 0))}) " - f"marker={components.get('marker', 0):.1f}(n={int(components.get('marker_count', 0))}) " - f"cassette={components.get('cassette', 0):.1f} " - f"prom={components.get('promoter', 0):.1f}(n={int(components.get('promoter_count', 0))}) " - f"term={components.get('terminator', 0):.1f}(n={int(components.get('terminator_count', 0))}) " - f"payload={components.get('payload', 0):.1f}(n={int(components.get('payload_count', 0))}) " - f"len_bonus={components.get('length_bonus', 0):.1f}({int(components.get('sequence_length', 0))}bp)" - ) - - if _REWARD_LOG_TIMINGS: - dt_ms = (time.perf_counter() - t0) * 1000.0 - logger.info( - f"reward.compute_score seq_len={len(sequence)} " - f"raw_score={raw_score:.2f} normalized={normalized_score:.4f} " - f"time_ms={dt_ms:.2f}" - ) - - return float(normalized_score) - - - - -# Backward compatibility alias -def get_plasmid_reward( - plasmid: Optional[str] = None, - *, - solution_str: Optional[str] = None, - data_source: Optional[str] = None, - ground_truth: Optional[str] = None, - extra_info: Optional[Dict] = None, - **kwargs, -) -> float: - """ - Backward compatibility wrapper. - Delegates to compute_score. - """ - seq = plasmid or solution_str or "" - return compute_score(data_source, seq, ground_truth, extra_info) diff --git a/tests/test_rewards.py b/tests/test_rewards.py index e250494..956ba29 100644 --- a/tests/test_rewards.py +++ b/tests/test_rewards.py @@ -1,9 +1,8 @@ import os import pytest -from src.rewards.scorer import Scorer -from src.rewards.reward_config import RewardConfig -from src.rewards import rewards as rewards_mod +from src.rewards.bioinformatics.scorer import Scorer +from src.rewards.bioinformatics.reward_config import RewardConfig def _read_fasta_sequence(path: str) -> str: @@ -77,22 +76,25 @@ def test_location_bonus_non_decreasing(): assert score_on >= score_off -def test_wrapper_score_sequence_bounds_and_consistency(): +def test_scorer_score_bounds(): data_dir = os.path.join(os.path.dirname(__file__), "data") path = os.path.join(data_dir, "pSC101.fasta") seq = _read_fasta_sequence(path) - # wrapper returns 0..100 - s100 = rewards_mod.score_sequence(seq) - print(f"testlog file={os.path.basename(path)} wrapper_score_0_100={s100:.2f}") - assert 0.0 <= s100 <= 100.0 + cfg = RewardConfig() + scorer = Scorer(cfg) + s, _ = scorer.score(seq) + print(f"testlog file={os.path.basename(path)} scorer_score_0_1={s:.4f}") + assert 0.0 <= s <= 1.0 -def test_score_completions_handles_empty(): +def test_simple_batch_handles_empty(): data_dir = os.path.join(os.path.dirname(__file__), "data") path = os.path.join(data_dir, "RF0G-IodoY.fasta") seq = _read_fasta_sequence(path) - scores = rewards_mod.score_completions([seq, ""]) - print(f"testlog file={os.path.basename(path)} completions_scores={scores}") - assert len(scores) == 2 - assert scores[1] == 0.0 - assert 0.0 <= scores[0] <= 100.0 + cfg = RewardConfig() + scorer = Scorer(cfg) + s_valid, _ = scorer.score(seq) + print(f"testlog file={os.path.basename(path)} simple_score_valid={s_valid}") + assert 0.0 <= s_valid <= 1.0 + with pytest.raises(Exception): + scorer.score("") From f159e0dd6430a35fbf29fb739c40baf596b56ea9 Mon Sep 17 00:00:00 2001 From: mcclain Date: Fri, 31 Oct 2025 15:57:02 +0000 Subject: [PATCH 06/11] fixing many small bugs --- src/runners/grpo.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/src/runners/grpo.py b/src/runners/grpo.py index 54b735f..b26e82f 100644 --- a/src/runners/grpo.py +++ b/src/runners/grpo.py @@ -9,6 +9,8 @@ import datetime from typing import List import wandb +from concurrent.futures import ThreadPoolExecutor +import re cfg = Config() MODEL_ID = cfg.model @@ -131,12 +133,32 @@ def select_prompt_and_extras(ds): reward_logger = RewardComponentLogger(log_frequency=10) -def batch_reward_fn(samples: List[str], **kwargs) -> List[float]: - rewards: List[float] = [] - for seq in samples: +def score_single(idx_and_seq): + idx, seq = idx_and_seq + try: score, components = scorer.score(seq) - rewards.append(float(score)) reward_logger.add_components(components, float(score)) + return float(score) + except Exception: + print(f"Warning: Failed to score completion {idx} (len={len(seq)}): {str(e)[:100]}") + return 0.0 + +def batch_reward_fn(prompts: List[str], completions: List[str], **kwargs) -> List[float]: + """ + TRL GRPO reward function signature. + Args: + prompts: List of prompt strings + completions: List of completion strings (without prompts) + Returns: + List of reward floats + """ + # Clean: remove spaces, uppercase, keep only valid DNA chars + cleaned = [re.sub(r'[^ATCG]', '', c.upper().replace(" ", "")) for c in completions] + + # Parallelize scoring + with ThreadPoolExecutor(max_workers=8) as executor: + rewards = list(executor.map(score_single, enumerate(cleaned))) + return rewards From 0b44694d0d8d5dbf7fc93958453121f610d438b2 Mon Sep 17 00:00:00 2001 From: mcclain Date: Fri, 31 Oct 2025 16:56:15 +0000 Subject: [PATCH 07/11] logging finally works --- src/rewards/bioinformatics/logger.py | 89 ++++-- src/rewards/bioinformatics/reward_config.py | 2 +- src/rewards/bioinformatics/scorer.py | 328 ++++++++++---------- src/runners/grpo.py | 36 ++- tests/test_rewards.py | 28 +- 5 files changed, 266 insertions(+), 217 deletions(-) diff --git a/src/rewards/bioinformatics/logger.py b/src/rewards/bioinformatics/logger.py index efa1904..6b6733a 100644 --- a/src/rewards/bioinformatics/logger.py +++ b/src/rewards/bioinformatics/logger.py @@ -29,45 +29,72 @@ def add_components(self, components: Dict[str, float], total_reward: float) -> N self.component_buffer["total_reward"].append(float(total_reward)) def on_step_end(self, args, state, control, **kwargs): - if self.log_frequency <= 0 or state.global_step == 0: - return - if (state.global_step % self.log_frequency) != 0: - return + """Called at the end of each training step to log reward components.""" + # Always try to log if there's data, regardless of step if not self.component_buffer["total_reward"]: return + + # Only log at specified frequency + if self.log_frequency > 0 and (state.global_step % self.log_frequency) != 0: + return log_dict: Dict[str, Any] = {} - try: - for name, values in self.component_buffer.items(): - if not values: - continue - arr = np.asarray(values, dtype=np.float32) - base = f"reward_components/{name}" - log_dict[f"{base}/mean"] = float(np.mean(arr)) - log_dict[f"{base}/std"] = float(np.std(arr)) - log_dict[f"{base}/min"] = float(np.min(arr)) - log_dict[f"{base}/max"] = float(np.max(arr)) + + # Log component statistics + for name, values in self.component_buffer.items(): + if not values: + continue + arr = np.asarray(values, dtype=np.float32) + base = f"reward_components/{name}" + log_dict[f"{base}/mean"] = float(np.mean(arr)) + log_dict[f"{base}/std"] = float(np.std(arr)) + log_dict[f"{base}/min"] = float(np.min(arr)) + log_dict[f"{base}/max"] = float(np.max(arr)) + + # Violation counters (fraction of sequences meeting criteria) + ori_vals = self.component_buffer.get("ori", []) + marker_vals = self.component_buffer.get("marker", []) + cds_vals = self.component_buffer.get("cds", []) + promoter_vals = self.component_buffer.get("promoter", []) + terminator_vals = self.component_buffer.get("terminator", []) + + if ori_vals: + log_dict["reward_violations/ori_perfect"] = float(sum(x >= 1.0 for x in ori_vals) / len(ori_vals)) + log_dict["reward_violations/ori_zero"] = float(sum(x == 0.0 for x in ori_vals) / len(ori_vals)) + if marker_vals: + log_dict["reward_violations/marker_perfect"] = float(sum(x >= 1.0 for x in marker_vals) / len(marker_vals)) + log_dict["reward_violations/marker_zero"] = float(sum(x == 0.0 for x in marker_vals) / len(marker_vals)) + if cds_vals: + log_dict["reward_violations/cds_perfect"] = float(sum(x >= 1.0 for x in cds_vals) / len(cds_vals)) + log_dict["reward_violations/cds_zero"] = float(sum(x == 0.0 for x in cds_vals) / len(cds_vals)) + if promoter_vals: + log_dict["reward_violations/promoter_perfect"] = float(sum(x >= 1.0 for x in promoter_vals) / len(promoter_vals)) + if terminator_vals: + log_dict["reward_violations/terminator_perfect"] = float(sum(x >= 1.0 for x in terminator_vals) / len(terminator_vals)) - # Simple violation counters - ori_vals = self.component_buffer.get("ori", []) - marker_vals = self.component_buffer.get("marker", []) - cds_vals = self.component_buffer.get("cds", []) - if ori_vals: - log_dict["reward_violations/ori_perfect"] = float(sum(x >= 1.0 for x in ori_vals) / len(ori_vals)) - if marker_vals: - log_dict["reward_violations/marker_perfect"] = float(sum(x >= 1.0 for x in marker_vals) / len(marker_vals)) - if cds_vals: - log_dict["reward_violations/cds_zero"] = float(sum(x == 0.0 for x in cds_vals) / len(cds_vals)) + # Log sample count + log_dict["reward_components/sample_count"] = len(self.component_buffer["total_reward"]) - # Periodic histograms - if (state.global_step % (self.log_frequency * 10)) == 0: - log_dict["reward_histograms/total_reward"] = wandb.Histogram(self.component_buffer["total_reward"]) # type: ignore[arg-type] + # Periodic histograms (every 10x log frequency) + if (state.global_step % (self.log_frequency * 10)) == 0 and len(self.component_buffer["total_reward"]) > 0: + log_dict["reward_histograms/total_reward"] = wandb.Histogram(self.component_buffer["total_reward"]) # type: ignore[arg-type] + if self.component_buffer["ori"]: log_dict["reward_histograms/ori"] = wandb.Histogram(self.component_buffer["ori"]) # type: ignore[arg-type] + if self.component_buffer["cds"]: log_dict["reward_histograms/cds"] = wandb.Histogram(self.component_buffer["cds"]) # type: ignore[arg-type] + if self.component_buffer["marker"]: + log_dict["reward_histograms/marker"] = wandb.Histogram(self.component_buffer["marker"]) # type: ignore[arg-type] - wandb.log(log_dict, step=int(state.global_step)) - finally: - for key in self.component_buffer: - self.component_buffer[key] = [] + # Log to wandb (don't specify step, let W&B use its own counter) + if log_dict: + try: + wandb.log(log_dict) + print(f"[RewardLogger] Logged {len(self.component_buffer['total_reward'])} samples at trainer step {state.global_step}") + except Exception as e: + print(f"[RewardLogger] Warning: Failed to log to wandb: {e}") + + # Clear buffers after logging + for key in self.component_buffer: + self.component_buffer[key] = [] diff --git a/src/rewards/bioinformatics/reward_config.py b/src/rewards/bioinformatics/reward_config.py index b91aecc..0ed3630 100644 --- a/src/rewards/bioinformatics/reward_config.py +++ b/src/rewards/bioinformatics/reward_config.py @@ -4,7 +4,7 @@ class RewardConfig(BaseModel): punish_mode: bool = True # penaize violations of the reward config as opposed to just not rewarding them - length_penalty: bool = True # penalize sequences that are too long or too short + length_penalty: bool = False # penalize sequences that are too long or too short min_length: Optional[int] = None max_length: Optional[int] = None location_aware: bool = True # reward sequences that are located in the correct location (e.g. prompoter then cds then terminatr) diff --git a/src/rewards/bioinformatics/scorer.py b/src/rewards/bioinformatics/scorer.py index 5b7b9b0..cf8fd3e 100644 --- a/src/rewards/bioinformatics/scorer.py +++ b/src/rewards/bioinformatics/scorer.py @@ -1,16 +1,39 @@ from src.rewards.bioinformatics.reward_config import RewardConfig import plasmidkit as pk -from typing import Any, List, Tuple, Union, Dict -from concurrent.futures import ThreadPoolExecutor +from typing import Any, List, Tuple, Dict import time -import hashlib -import os + + +class _Feat: + """Simple feature container for annotation merging.""" + def __init__(self, type: str, id: str | None, start: int, end: int, strand: str | None, evidence: Any = None): + self.type = type + self.id = id + self.start = int(start) + self.end = int(end) + self.strand = strand or "+" + self.evidence = evidence or {} class Scorer: + """ + Scores plasmid sequences based on biological features (ori, promoter, terminator, marker, CDS). + + Uses plasmidkit for annotation, then computes weighted scores for each component. + Optionally applies length penalties and location-aware bonuses for gene cassettes. + """ def __init__(self, reward_config: RewardConfig): self.reward_config = reward_config - # Fixed set for simplicity; location/length are folded into component scores/final modulation + if self.reward_config.length_penalty: + assert ( + self.reward_config.min_length is not None or self.reward_config.max_length is not None + ), "At least one of min_length or max_length must be set if length_penalty is True" + if ( + self.reward_config.min_length is not None + and self.reward_config.max_length is not None + ): + assert self.reward_config.min_length < self.reward_config.max_length, "min_length must be less than max_length" + self.score_functions = [ self.score_ori, self.score_promoter, @@ -26,58 +49,49 @@ def __init__(self, reward_config: RewardConfig): self.reward_config.cds_weight, ] total_weight = sum(self.weights) - if total_weight > 0: - self.weights = [w / total_weight for w in self.weights] - else: - uniform = 1.0 / len(self.score_functions) - self.weights = [uniform for _ in self.score_functions] + assert total_weight > 0, "Total weight must be greater than 0" + self.weights = [w / total_weight for w in self.weights] def annotate(self, sequence: str) -> Any: + """Annotate sequence with plasmidkit and merge overlapping features.""" + assert sequence, "sequence cannot be empty" raw = pk.annotate(sequence, is_sequence=True) return self._preprocess_annotations(raw) - # --- preprocessing helpers --- - class _Feat: - def __init__(self, type: str, id: str | None, start: int, end: int, strand: str | None, evidence: Any = None): - self.type = type - self.id = id - self.start = int(start) - self.end = int(end) - self.strand = strand or "+" - self.evidence = evidence or {} - - @staticmethod - def _len(x: Any) -> int: - return max(0, int(getattr(x, "end", 0)) - int(getattr(x, "start", 0))) - @staticmethod def _overlap_len(a: Any, b: Any) -> int: - s1, e1 = int(getattr(a, "start", 0)), int(getattr(a, "end", 0)) - s2, e2 = int(getattr(b, "start", 0)), int(getattr(b, "end", 0)) + """Calculate overlap length between two features.""" + s1, e1 = int(a.start), int(a.end) + s2, e2 = int(b.start), int(b.end) lo = max(min(s1, e1), min(s2, e2)) hi = min(max(s1, e1), max(s2, e2)) return max(0, hi - lo) - def _to_feat(self, x: Any) -> "Scorer._Feat": - return Scorer._Feat( - type=(getattr(x, "type", None) or "").lower(), - id=getattr(x, "id", None), - start=int(getattr(x, "start", 0)), - end=int(getattr(x, "end", 0)), - strand=getattr(x, "strand", "+"), - evidence=getattr(x, "evidence", {}), + def _to_feat(self, x: Any) -> "_Feat": + """Convert annotation object to internal _Feat representation.""" + return _Feat( + type=x.type.lower() if x.type else "", + id=x.id if hasattr(x, "id") else None, + start=int(x.start), + end=int(x.end), + strand=x.strand if hasattr(x, "strand") else "+", + evidence=x.evidence if hasattr(x, "evidence") else {}, ) - def _merge_group(self, feats: List[Any], threshold: float, *, respect_strand: bool) -> List["Scorer._Feat"]: + def _merge_group(self, feats: List[Any], threshold: float, *, respect_strand: bool) -> List["_Feat"]: + """Merge overlapping features of the same type based on overlap threshold.""" if not feats: return [] items = [self._to_feat(f) for f in feats] items.sort(key=lambda f: (f.strand, f.start, f.end)) - merged: List[Scorer._Feat] = [] + merged: List[_Feat] = [] cur = items[0] for nxt in items[1:]: ovl = self._overlap_len(cur, nxt) - min_len = max(1, min(self._len(cur), self._len(nxt))) + # Inline length calculation + cur_len = max(0, cur.end - cur.start) + nxt_len = max(0, nxt.end - nxt.start) + min_len = max(1, min(cur_len, nxt_len)) strands_compatible = (cur.strand == nxt.strand) or (not respect_strand) if ovl / float(min_len) >= threshold and strands_compatible: cur.start = min(cur.start, nxt.start) @@ -90,36 +104,42 @@ def _merge_group(self, feats: List[Any], threshold: float, *, respect_strand: bo return merged def _preprocess_annotations(self, annotations: Any) -> List[Any]: + """ + Merge overlapping annotations and filter out CDS overlapping with other feature types. + + Groups features by type, merges overlapping ones based on threshold, and removes + CDS annotations that overlap with ori/promoter/terminator/marker features. + """ feats = list(annotations) thr = float(self.reward_config.overlap_merge_threshold) - type_key = lambda x: (getattr(x, "type", None) or "").lower() + type_key = lambda x: x.type.lower() if x.type else "" - # collect groups + # Collect groups by type groups: Dict[str, List[Any]] = {} for f in feats: groups.setdefault(type_key(f), []).append(f) - # merge per group for relevant types - merged_groups: Dict[str, List[Scorer._Feat]] = {} + # Merge per group for relevant types + merged_groups: Dict[str, List[_Feat]] = {} for t in ("rep_origin", "ori", "origin_of_replication", "promoter", "terminator", "marker", "cds"): if t in groups: # Ignore strand for ORI and marker; respect for others respect = t not in ("rep_origin", "ori", "origin_of_replication", "marker") merged_groups[t] = self._merge_group(groups[t], thr, respect_strand=respect) - # suppress CDS if overlaps any non-CDS (ori/promoter/terminator/marker) - non_cds: List[Scorer._Feat] = [] + # Suppress CDS if overlaps any non-CDS (ori/promoter/terminator/marker) + non_cds: List[_Feat] = [] for t in ("rep_origin", "ori", "origin_of_replication", "promoter", "terminator", "marker"): non_cds.extend(merged_groups.get(t, [])) - filtered_cds: List[Scorer._Feat] = [] + filtered_cds: List[_Feat] = [] for c in merged_groups.get("cds", []): if any(self._overlap_len(c, o) > 0 for o in non_cds): continue filtered_cds.append(c) merged_groups["cds"] = filtered_cds - # rebuild final list: prefer merged groups for those types; keep other features as-is + # Rebuild final list: prefer merged groups for those types; keep other features as-is final: List[Any] = [] merged_types = set(merged_groups.keys()) for t, items in merged_groups.items(): @@ -130,94 +150,49 @@ def _preprocess_annotations(self, annotations: Any) -> List[Any]: final.append(f) return final - @staticmethod - def _read_fasta_file(path: str) -> str: - with open(path, "r") as f: - lines = [ln.strip() for ln in f.readlines()] - seq_lines = [ln for ln in lines if ln and not ln.startswith(">")] - return "".join(seq_lines).replace(" ", "").upper() - @staticmethod - def _derive_source(sequence: str, provided: str | None) -> str: - if provided: - return provided - # Try FASTA header if present - if sequence and sequence.lstrip().startswith(">"): - first = sequence.splitlines()[0].lstrip()[1:].strip() - return first.split()[0] if first else "" - # Fallback: length + sha1 fingerprint - sha = hashlib.sha1((sequence or "").encode("utf-8")).hexdigest()[:8] - return f"seq:bp{len(sequence or '')}:sha{sha}" - - def runner(self, sequence: str, annotations: Any) -> List[float]: - """ - Takes a list of runnable functions with the same signature and runs them in parallel. - - Args: - sequence: The sequence to score - annotations: The annotations to score - - Returns: - A list of scores in the same order as `self.score_functions`. - """ - with ThreadPoolExecutor(max_workers=len(self.score_functions)) as executor: - futures = [executor.submit(func, sequence, annotations) for func in self.score_functions] - return [future.result() for future in futures] - - # --- helpers --- - @staticmethod - def _feat_type(x: Any) -> str: - return (getattr(x, "type", None) or "").lower() - - @staticmethod - def _feat_id(x: Any) -> str: - return (getattr(x, "id", None) or "").lower() - - @staticmethod - def _strand(x: Any) -> str: - return getattr(x, "strand", "+") - - @staticmethod - def _start(x: Any) -> int: - return int(getattr(x, "start", 0)) - - @staticmethod - def _end(x: Any) -> int: - return int(getattr(x, "end", 0)) - @staticmethod def _distance(a: Any, b: Any) -> int: - if Scorer._end(a) < Scorer._start(b): - return Scorer._start(b) - Scorer._end(a) - if Scorer._end(b) < Scorer._start(a): - return Scorer._start(a) - Scorer._end(b) + """Calculate distance between two non-overlapping features (0 if overlapping).""" + a_start, a_end = int(a.start), int(a.end) + b_start, b_end = int(b.start), int(b.end) + if a_end < b_start: + return b_start - a_end + if b_end < a_start: + return a_start - b_end return 0 @staticmethod def _in_order_same_strand(a: Any, b: Any) -> bool: - if Scorer._strand(a) != Scorer._strand(b): + """Check if features a and b are on same strand and in correct order.""" + strand_a = a.strand if hasattr(a, "strand") else "+" + strand_b = b.strand if hasattr(b, "strand") else "+" + if strand_a != strand_b: return False - if Scorer._strand(a) == "+": - return Scorer._start(a) <= Scorer._start(b) - return Scorer._end(a) >= Scorer._end(b) - - def _prox_points(self, d: int) -> float: - return ( - float(self.reward_config.cassette_proximity_points) - if d <= int(self.reward_config.proximity_threshold_bp) - else 0.0 - ) + start_a = int(a.start) + start_b = int(b.start) + if strand_a == "+": + return start_a <= start_b + end_a = int(a.end) + end_b = int(b.end) + return end_a >= end_b @staticmethod def _filter_allowed(xs: List[Any], allowed: List[str] | None) -> List[Any]: + """Filter features to only those matching allowed IDs.""" if not allowed: return xs allowed_l = [a.lower() for a in allowed] - return [x for x in xs if any(a in Scorer._feat_id(x) for a in allowed_l)] + return [x for x in xs if hasattr(x, "id") and x.id and any(a in x.id.lower() for a in allowed_l)] def _count_score(self, count: int, min_req: int, max_allowed: int) -> float: + """ + Score feature count based on min/max requirements. + + Returns 1.0 if within range, proportional score if below min, and penalty if above max. + """ if max_allowed <= 0: return 0.0 - # Below minimum: proportion of requirement; softer if punish_mode + # Below minimum: proportion of requirement if count < min_req: proportion = (count / float(min_req)) if min_req > 0 else 0.0 return max(0.0, min(1.0, proportion * (0.5 if self.reward_config.punish_mode else 1.0))) @@ -228,108 +203,150 @@ def _count_score(self, count: int, min_req: int, max_allowed: int) -> float: return 1.0 def _compute_cassette_points(self, promoters: List[Any], cdss: List[Any], terminators: List[Any]) -> float: - # Returns summed raw points across top-N cassettes + """ + Compute bonus points for well-formed gene cassettes (promoter -> CDS -> terminator). + + Awards points for correct order and proximity on same strand. + Returns summed raw points across top-N cassettes. + """ triples: List[Tuple[Any, Any, Any, float]] = [] for p in promoters: - cds_same = [c for c in cdss if self._strand(c) == self._strand(p)] + # Find closest CDS on same strand + p_strand = p.strand if hasattr(p, "strand") else "+" + cds_same = [c for c in cdss if (c.strand if hasattr(c, "strand") else "+") == p_strand] if not cds_same: continue cds_same.sort(key=lambda c: self._distance(p, c)) c = cds_same[0] - term_cands = [t for t in terminators if self._strand(t) == self._strand(c)] + + # Find closest terminator on same strand as CDS + c_strand = c.strand if hasattr(c, "strand") else "+" + term_cands = [t for t in terminators if (t.strand if hasattr(t, "strand") else "+") == c_strand] term_cands.sort(key=lambda t: self._distance(c, t)) t = term_cands[0] if term_cands else None + pts = 0.0 - # promoter -> CDS + # Promoter -> CDS + dist_pc = self._distance(p, c) if self._in_order_same_strand(p, c): pts += self.reward_config.cassette_order_points - pts += min(self.reward_config.cassette_proximity_points, self._prox_points(self._distance(p, c))) + if dist_pc <= self.reward_config.proximity_threshold_bp: + pts += self.reward_config.cassette_proximity_points else: - pts += min(3.0, self._prox_points(self._distance(p, c))) - # CDS -> terminator + if dist_pc <= self.reward_config.proximity_threshold_bp: + pts += min(3.0, self.reward_config.cassette_proximity_points) + + # CDS -> Terminator if t is not None: + dist_ct = self._distance(c, t) if self._in_order_same_strand(c, t): pts += self.reward_config.cassette_order_points - pts += min(self.reward_config.cassette_proximity_points, self._prox_points(self._distance(c, t))) + if dist_ct <= self.reward_config.proximity_threshold_bp: + pts += self.reward_config.cassette_proximity_points else: - pts += min(3.0, self._prox_points(self._distance(c, t))) + if dist_ct <= self.reward_config.proximity_threshold_bp: + pts += min(3.0, self.reward_config.cassette_proximity_points) + triples.append((p, c, t, pts)) + + # Take top N cassettes triples.sort(key=lambda x: x[3], reverse=True) top = triples[: int(self.reward_config.cassette_max_cassettes)] - total = 0.0 - for (_, _, _, pts) in top: - total += min(self.reward_config.cassette_max_points_per, float(pts)) - return total + return sum(min(self.reward_config.cassette_max_points_per, pts) for _, _, _, pts in top) - # --- component scores (0..1) --- def score_ori(self, seq: str, annotations: Any) -> float: + """Score origin of replication features.""" feats = list(annotations) - oris = [x for x in feats if self._feat_type(x) in ("rep_origin", "ori", "origin_of_replication")] + oris = [x for x in feats if x.type and x.type.lower() in ("rep_origin", "ori", "origin_of_replication")] oris = self._filter_allowed(oris, self.reward_config.allowed_oris) return self._count_score(len(oris), self.reward_config.ori_min, self.reward_config.ori_max) def score_promoter(self, seq: str, annotations: Any) -> float: + """Score promoter features.""" feats = list(annotations) - promoters = [x for x in feats if self._feat_type(x) == "promoter"] + promoters = [x for x in feats if x.type and x.type.lower() == "promoter"] promoters = self._filter_allowed(promoters, self.reward_config.allowed_promoters) - score = self._count_score(len(promoters), self.reward_config.promoter_min, self.reward_config.promoter_max) - # small standalone credit (capped implicitly by count scoring) - return score + return self._count_score(len(promoters), self.reward_config.promoter_min, self.reward_config.promoter_max) def score_terminator(self, seq: str, annotations: Any) -> float: + """Score terminator features.""" feats = list(annotations) - terms = [x for x in feats if self._feat_type(x) == "terminator"] + terms = [x for x in feats if x.type and x.type.lower() == "terminator"] terms = self._filter_allowed(terms, self.reward_config.allowed_terminators) return self._count_score(len(terms), self.reward_config.terminator_min, self.reward_config.terminator_max) def score_marker(self, seq: str, annotations: Any) -> float: + """Score selectable marker features (binary: present or absent).""" feats = list(annotations) - markers = [x for x in feats if self._feat_type(x) == "marker"] + markers = [x for x in feats if x.type and x.type.lower() == "marker"] markers = self._filter_allowed(markers, self.reward_config.allowed_markers) if markers: return 1.0 return 0.0 if self.reward_config.punish_mode else 0.5 def score_cds(self, seq: str, annotations: Any) -> float: + """Score CDS features with optional location-aware cassette bonus.""" feats = list(annotations) - cdss = [x for x in feats if self._feat_type(x) == "cds"] + cdss = [x for x in feats if x.type and x.type.lower() == "cds"] cdss = self._filter_allowed(cdss, self.reward_config.allowed_cds) base = self._count_score(len(cdss), self.reward_config.cds_min, self.reward_config.cds_max) if not self.reward_config.location_aware: return base - promoters = [x for x in feats if self._feat_type(x) == "promoter"] - terminators = [x for x in feats if self._feat_type(x) == "terminator"] + # Add cassette bonus for properly arranged promoter->CDS->terminator + promoters = [x for x in feats if x.type and x.type.lower() == "promoter"] + terminators = [x for x in feats if x.type and x.type.lower() == "terminator"] cassette_pts = self._compute_cassette_points(promoters, cdss, terminators) max_total = self.reward_config.cassette_max_cassettes * self.reward_config.cassette_max_points_per - bonus_scale = float(self.reward_config.location_bonus_scale) - cassette_bonus = bonus_scale * (cassette_pts / max_total if max_total > 0 else 0.0) + cassette_bonus = self.reward_config.location_bonus_scale * (cassette_pts / max_total if max_total > 0 else 0.0) return max(0.0, min(1.0, base + cassette_bonus)) - def _length_factor(self, seq: str) -> float: + def score_length(self, seq: str, annotations: Any) -> float: + """Score sequence length (penalty if outside allowed range).""" if not self.reward_config.length_penalty: return 1.0 - L = len(seq or "") + L = len(seq) mn = self.reward_config.min_length mx = self.reward_config.max_length - if mn is None and mx is None: - return 1.0 - # inside range: full credit; outside: halve if punish_mode else no change + assert mn is not None or mx is not None, "min_length or max_length must be set if length_penalty is True" + # Inside range: full credit; outside: penalty if punish_mode in_range = (mn is None or L >= mn) and (mx is None or L <= mx) if in_range: return 1.0 return float(self.reward_config.violation_penalty_factor) if self.reward_config.punish_mode else 1.0 def score(self, seq: str, source: str | None = None) -> Tuple[float, Dict[str, float]]: + """ + Score a DNA sequence based on biological features. + + Args: + seq: DNA sequence to score + source: Optional identifier for logging + + Returns: + Tuple of (final_score, component_scores_dict) + """ + assert seq, "sequence cannot be empty" t0 = time.perf_counter() - src = Scorer._derive_source(seq, source) + src = source if source else f"seq_{len(seq)}bp" + + # Annotate and score each component annotations = self.annotate(seq) - results = self.runner(seq, annotations) - ori, prom, term, mark, cds = results + ori = self.score_ori(seq, annotations) + prom = self.score_promoter(seq, annotations) + term = self.score_terminator(seq, annotations) + mark = self.score_marker(seq, annotations) + cds = self.score_cds(seq, annotations) + length_factor = self.score_length(seq, annotations) + + # Weighted sum + results = [ori, prom, term, mark, cds] base = sum(w * r for w, r in zip(self.weights, results)) - length_factor = self._length_factor(seq) + + # Apply length penalty as multiplier final = max(0.0, min(1.0, base * length_factor)) + components: Dict[str, float] = { "ori": float(ori), "promoter": float(prom), @@ -338,11 +355,6 @@ def score(self, seq: str, source: str | None = None) -> Tuple[float, Dict[str, f "cds": float(cds), "length_factor": float(length_factor), } + dt_ms = (time.perf_counter() - t0) * 1000.0 - cfg = self.reward_config.model_dump(exclude_none=True) - print(f"source={src} config={cfg} score={final:.4f} len={len(seq)} time_ms={dt_ms:.2f} components={components}") - return float(final), components - - def score_fasta(self, fasta_path: str) -> Tuple[float, Dict[str, float]]: - seq = Scorer._read_fasta_file(fasta_path) - return self.score(seq, source=os.path.basename(fasta_path)) \ No newline at end of file + return float(final), components \ No newline at end of file diff --git a/src/runners/grpo.py b/src/runners/grpo.py index b26e82f..916031b 100644 --- a/src/runners/grpo.py +++ b/src/runners/grpo.py @@ -10,6 +10,7 @@ from typing import List import wandb from concurrent.futures import ThreadPoolExecutor +from threading import Lock import re cfg = Config() @@ -35,13 +36,10 @@ def select_prompt_and_extras(ds): train_ds = load_dataset("parquet", data_files=TRAIN_PARQUET, split="train") train_ds = select_prompt_and_extras(train_ds) -try: - eval_ds = load_dataset("parquet", data_files=VAL_PARQUET, split="train") - eval_ds = select_prompt_and_extras(eval_ds) - use_eval = True -except Exception: - eval_ds = None - use_eval = False +eval_ds = load_dataset("parquet", data_files=VAL_PARQUET, split="train") +eval_ds = select_prompt_and_extras(eval_ds) +use_eval = True + tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True, trust_remote_code=True) @@ -118,6 +116,10 @@ def select_prompt_and_extras(ds): # ---- reward config and scorer ---- reward_config = RewardConfig( + punish_mode=False, + length_penalty=True, + min_length=1000, + max_length=30000, ori_min=1, ori_max=1, promoter_min=1, @@ -130,18 +132,23 @@ def select_prompt_and_extras(ds): scorer = Scorer(reward_config) -reward_logger = RewardComponentLogger(log_frequency=10) +reward_logger = RewardComponentLogger(log_frequency=1) + +# Thread-safe storage for component data +component_lock = Lock() def score_single(idx_and_seq): idx, seq = idx_and_seq try: score, components = scorer.score(seq) - reward_logger.add_components(components, float(score)) - return float(score) - except Exception: + # Thread-safe addition to logger + with component_lock: + reward_logger.add_components(components, float(score)) + return float(score), components + except Exception as e: print(f"Warning: Failed to score completion {idx} (len={len(seq)}): {str(e)[:100]}") - return 0.0 + return 0.0, None def batch_reward_fn(prompts: List[str], completions: List[str], **kwargs) -> List[float]: """ @@ -157,7 +164,10 @@ def batch_reward_fn(prompts: List[str], completions: List[str], **kwargs) -> Lis # Parallelize scoring with ThreadPoolExecutor(max_workers=8) as executor: - rewards = list(executor.map(score_single, enumerate(cleaned))) + results = list(executor.map(score_single, enumerate(cleaned))) + + # Extract just the rewards for return + rewards = [r[0] for r in results] return rewards diff --git a/tests/test_rewards.py b/tests/test_rewards.py index 956ba29..de98c65 100644 --- a/tests/test_rewards.py +++ b/tests/test_rewards.py @@ -23,19 +23,19 @@ def test_scorer_components_and_bounds(): scorer = Scorer(cfg) for p in fasta_paths: - score, components = scorer.score_fasta(p) - cfg_dump = cfg.model_dump(exclude_none=True) seq = _read_fasta_sequence(p) + score, components = scorer.score(seq, source=os.path.basename(p)) + cfg_dump = cfg.model_dump(exclude_none=True) ann = scorer.annotate(seq) ann_view = [ { - "type": getattr(a, "type", None), - "id": getattr(a, "id", None), - "start": getattr(a, "start", None), - "end": getattr(a, "end", None), - "strand": getattr(a, "strand", None), + "type": a.type if hasattr(a, "type") else None, + "id": a.id if hasattr(a, "id") else None, + "start": a.start if hasattr(a, "start") else None, + "end": a.end if hasattr(a, "end") else None, + "strand": a.strand if hasattr(a, "strand") else None, } - for a in ann if getattr(a, "type", None) != "restriction_site" + for a in ann if hasattr(a, "type") and a.type != "restriction_site" ] print( f"testlog file={os.path.basename(p)} config={cfg_dump} score={score:.4f} components={components} annotations={ann_view}" @@ -61,13 +61,13 @@ def test_location_bonus_non_decreasing(): ann = scorer_on.annotate(seq) ann_view = [ { - "type": getattr(a, "type", None), - "id": getattr(a, "id", None), - "start": getattr(a, "start", None), - "end": getattr(a, "end", None), - "strand": getattr(a, "strand", None), + "type": a.type if hasattr(a, "type") else None, + "id": a.id if hasattr(a, "id") else None, + "start": a.start if hasattr(a, "start") else None, + "end": a.end if hasattr(a, "end") else None, + "strand": a.strand if hasattr(a, "strand") else None, } - for a in ann if getattr(a, "type", None) != "restriction_site" + for a in ann if hasattr(a, "type") and a.type != "restriction_site" ] print(f"testlog file={os.path.basename(path)} config_off={cfg_off.model_dump()} score_off={score_off:.4f} comp_off={comp_off}") print(f"testlog file={os.path.basename(path)} config_on={cfg_on.model_dump()} score_on={score_on:.4f} comp_on={comp_on} annotations={ann_view}") From e1756d7d475976eed3a4c274058a53f78c32aa98 Mon Sep 17 00:00:00 2001 From: mcclain Date: Fri, 31 Oct 2025 17:12:54 +0000 Subject: [PATCH 08/11] big sweep --- SWEEPS.md | 96 +++++++++++++++++++ docker-compose.yaml | 18 ++++ src/runners/grpo.py | 158 +++++++++++++++---------------- src/runners/grpo_sweep.py | 191 ++++++++++++++++++++++++++++++++++++++ sweep_config.yaml | 107 +++++++++++++++++++++ 5 files changed, 488 insertions(+), 82 deletions(-) create mode 100644 SWEEPS.md create mode 100644 src/runners/grpo_sweep.py create mode 100644 sweep_config.yaml diff --git a/SWEEPS.md b/SWEEPS.md new file mode 100644 index 0000000..cfaf340 --- /dev/null +++ b/SWEEPS.md @@ -0,0 +1,96 @@ +# Hyperparameter Sweeps with W&B + +This directory contains configuration for running hyperparameter sweeps using Weights & Biases. + +## Quick Start + +### 1. Initialize a Sweep + +```bash +wandb sweep sweep_config.yaml +``` + +This will output a sweep ID like: `mcclain/plasmidrl-grpo-sweeps/abc123xyz` + +### 2. Run Sweep Agent(s) + +**Single agent:** +```bash +SWEEP_ID=mcclain/plasmidrl-grpo-sweeps/abc123xyz docker compose up grpo-sweep +``` + +**Multiple parallel agents (recommended for faster sweeps):** +```bash +SWEEP_ID=mcclain/plasmidrl-grpo-sweeps/abc123xyz docker compose up --scale grpo-sweep=3 +``` + +### 3. Monitor Progress + +Visit your W&B dashboard: +- https://wandb.ai/mcclain/plasmidrl-grpo-sweeps + +## Sweep Configuration + +The sweep is configured in `sweep_config.yaml` with the following parameters: + +### Training Hyperparameters +- `learning_rate`: 1e-6 to 1e-4 (log uniform) +- `per_device_train_batch_size`: 8, 16, or 32 +- `num_generations`: 4, 8, or 16 +- `temperature`: 0.7 to 1.4 +- `top_p`: 0.85 to 0.95 + +### GRPO Parameters +- `beta`: 1e-4 to 1e-2 (KL penalty coefficient) +- `epsilon`: 0.1 to 0.3 (PPO-style clipping) + +### Reward Configuration +- **Constraints**: min/max lengths, counts for each component +- **Weights**: Can be set to 0.0 to disable a component + - `reward_ori_weight`: 0.0, 0.5, 1.0, or 2.0 + - `reward_promoter_weight`: 0.0, 0.5, 1.0, or 2.0 + - `reward_terminator_weight`: 0.0, 0.25, 0.5, or 1.0 + - `reward_marker_weight`: 0.0, 0.5, 1.0, or 2.0 + - `reward_cds_weight`: 0.0, 0.5, 1.0, or 2.0 +- **Penalties**: `reward_punish_mode`, `reward_length_penalty` +- **Features**: `reward_location_aware` (cassette bonuses) + +## Optimization Strategy + +The sweep uses **Bayesian optimization** to intelligently explore the hyperparameter space: +- **Maximizes**: `reward_components/total_reward/mean` - mean reward across all training steps +- **Quick evaluation**: Each run completes in ~100 steps for fast iteration +- **Early termination**: Hyperband stops poorly performing runs after 20 steps +- **Smart search**: Focuses on promising hyperparameter regions + +## Customizing Sweeps + +Edit `sweep_config.yaml` to: +1. Change the search method (`bayes`, `grid`, `random`) +2. Add/remove parameters +3. Adjust parameter ranges +4. Change the optimization metric +5. Modify early termination settings + +## Best Practices + +1. **Start small**: Run a few trials manually to validate config +2. **Use multiple agents**: Parallel agents speed up sweeps significantly +3. **Monitor early**: Check results after 5-10 runs to ensure valid config +4. **Stop bad sweeps**: If all runs fail, stop and fix the config +5. **Save good configs**: Note the best hyperparameters for future runs + +## Troubleshooting + +**Sweep not starting:** +- Verify W&B login: `wandb login` +- Check SWEEP_ID format includes entity/project/id + +**All runs failing:** +- Check logs: `docker compose logs grpo-sweep` +- Verify reward config constraints are achievable + +**Out of memory:** +- Reduce `per_device_train_batch_size` in sweep config +- Reduce `num_generations` parameter range + diff --git a/docker-compose.yaml b/docker-compose.yaml index 00ca6c1..edc2a03 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -62,6 +62,24 @@ services: - WANDB_TAGS=["plasmid","rl","trl","es"] - WANDB_NOTES=TRL Evolution Strategies optimization + # W&B Sweep Agent for hyperparameter optimization + grpo-sweep: + <<: *common + volumes: + - ./:/mcclain + - /mnt/s3/phd-research-storage-1758274488:/s3:rw + command: > + bash -lc " + mkdir -p /tmp/wandb && \ + uv sync --frozen && \ + wandb agent ${SWEEP_ID} + " + environment: + - WANDB_ENTITY=mcclain + - WANDB_PROJECT=plasmidrl-grpo-sweeps + - WANDB_DIR=/tmp/wandb + - SWEEP_ID=${SWEEP_ID} + # VERL: PPO trainer verl-ppo: build: diff --git a/src/runners/grpo.py b/src/runners/grpo.py index 916031b..9184df0 100644 --- a/src/runners/grpo.py +++ b/src/runners/grpo.py @@ -1,5 +1,5 @@ from datasets import load_dataset -from transformers import AutoTokenizer, set_seed, TrainerCallback +from transformers import AutoTokenizer from trl import GRPOTrainer, GRPOConfig import torch from src.config import Config @@ -13,49 +13,38 @@ from threading import Lock import re +# Configuration cfg = Config() -MODEL_ID = cfg.model -TRAIN_PARQUET = cfg.train_dataset -VAL_PARQUET = cfg.val_dataset - -run_name = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") -run_name = f"grpo-{run_name}" -save_path = f"/s3/checkpoints/verl-grpo/{run_name}" -SEED = 42 - -# ---- dataset (keep your extra cols for reward) ---- -PROMPT_KEY = "prompt" -KEEP_EXTRA_COLS = ["data_source", "ability", "reward_model", "extra_info"] - - -def select_prompt_and_extras(ds): - cols = set(ds.column_names) - keep = ["prompt"] + [c for c in KEEP_EXTRA_COLS if c in cols] - return ds.select_columns(keep) - -train_ds = load_dataset("parquet", data_files=TRAIN_PARQUET, split="train") -train_ds = select_prompt_and_extras(train_ds) - -eval_ds = load_dataset("parquet", data_files=VAL_PARQUET, split="train") -eval_ds = select_prompt_and_extras(eval_ds) -use_eval = True - +run_name = f"grpo-{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" + +# Dataset loading +def load_train_val_datasets(): + """Load and preprocess training and validation datasets.""" + def select_prompt_column(ds): + cols = set(ds.column_names) + keep_cols = ["prompt"] + [c for c in ["data_source", "ability", "reward_model", "extra_info"] if c in cols] + return ds.select_columns(keep_cols) + + train_ds = load_dataset("parquet", data_files=cfg.train_dataset, split="train") + val_ds = load_dataset("parquet", data_files=cfg.val_dataset, split="train") + + return select_prompt_column(train_ds), select_prompt_column(val_ds) +train_ds, eval_ds = load_train_val_datasets() -tok = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True, trust_remote_code=True) +# Tokenizer setup +tok = AutoTokenizer.from_pretrained(cfg.model, use_fast=True, trust_remote_code=True) tok.padding_side = "left" - -# Remap specials to the correct strings (these already exist in your vocab) tok.eos_token = "" tok.bos_token = "" tok.pad_token = "[PAD]" -# Re-assert -assert tok.eos_token_id == 30001, tok.eos_token_id -assert tok.bos_token_id == 30000, tok.bos_token_id -assert tok.pad_token_id == 3, tok.pad_token_id +# Validate token IDs +assert tok.eos_token_id == 30001, f"Expected eos_token_id=30001, got {tok.eos_token_id}" +assert tok.bos_token_id == 30000, f"Expected bos_token_id=30000, got {tok.bos_token_id}" +assert tok.pad_token_id == 3, f"Expected pad_token_id=3, got {tok.pad_token_id}" -# Pass IDs explicitly to the model so nothing “helpfully” changes at runtime +# Model initialization kwargs model_init_kwargs = { "trust_remote_code": True, "eos_token_id": tok.eos_token_id, @@ -63,12 +52,12 @@ def select_prompt_and_extras(ds): "pad_token_id": tok.pad_token_id, } - -# ---- GRPO config (keys verified against TRL docs) ---- +# Training configuration args = GRPOConfig( model_init_kwargs=model_init_kwargs, - # transformers-style - output_dir=save_path, + output_dir=f"/s3/checkpoints/verl-grpo/{run_name}", + + # Training parameters num_train_epochs=20, learning_rate=3e-6, lr_scheduler_type="constant", @@ -76,45 +65,48 @@ def select_prompt_and_extras(ds): per_device_train_batch_size=16, gradient_accumulation_steps=1, max_steps=-1, + max_grad_norm=0.5, + seed=42, + + # Logging and checkpointing save_strategy="steps", save_steps=100, logging_strategy="steps", logging_steps=1, report_to=["wandb"], + + # Evaluation + do_eval=True, + eval_strategy="steps", + eval_steps=100, + + # Optimization bf16=torch.cuda.is_available(), gradient_checkpointing=False, - max_grad_norm=0.5, - seed=SEED, - do_eval=use_eval, - eval_strategy="steps" if use_eval else "no", - eval_steps=100 if use_eval else None, - - # model/ref-model handling - disable_dropout=True, # stabilizes ref-policy logprobs - - # data & generation - remove_unused_columns=False, # keep your extras for reward_fn + + # GRPO-specific + beta=1e-3, + epsilon=0.2, + loss_type="bnpo", + scale_rewards=True, + mask_truncated_completions=False, + disable_dropout=True, + + # Generation parameters + remove_unused_columns=False, max_prompt_length=1024, - num_generations=8, + num_generations=8, max_completion_length=256, - #min_completion_length=16, #not supported currently temperature=0.95, top_p=0.90, - - # GRPO specifics - beta=1e-3, # KL in loss → auto-loads ref model - epsilon=0.2, # PPO-style clip (replaces cliprange) - loss_type="bnpo", # token-level normalization; avoids length bias - scale_rewards=True, - mask_truncated_completions=False, - - # vLLM (colocated serverless flag is not a key here) + + # vLLM configuration use_vllm=True, vllm_gpu_memory_utilization=0.15, - vllm_mode="colocate" + vllm_mode="colocate", ) -# ---- reward config and scorer ---- +# Reward configuration reward_config = RewardConfig( punish_mode=False, length_penalty=True, @@ -130,19 +122,17 @@ def select_prompt_and_extras(ds): marker_max=2, ) +# Initialize scorer and logger scorer = Scorer(reward_config) - reward_logger = RewardComponentLogger(log_frequency=1) - -# Thread-safe storage for component data component_lock = Lock() - +# Reward function def score_single(idx_and_seq): + """Score a single sequence and log components thread-safely.""" idx, seq = idx_and_seq try: score, components = scorer.score(seq) - # Thread-safe addition to logger with component_lock: reward_logger.add_components(components, float(score)) return float(score), components @@ -152,35 +142,38 @@ def score_single(idx_and_seq): def batch_reward_fn(prompts: List[str], completions: List[str], **kwargs) -> List[float]: """ - TRL GRPO reward function signature. + Compute rewards for a batch of completions. + Args: prompts: List of prompt strings completions: List of completion strings (without prompts) + Returns: - List of reward floats + List of reward scores """ - # Clean: remove spaces, uppercase, keep only valid DNA chars + # Clean sequences: remove non-DNA characters cleaned = [re.sub(r'[^ATCG]', '', c.upper().replace(" ", "")) for c in completions] # Parallelize scoring with ThreadPoolExecutor(max_workers=8) as executor: results = list(executor.map(score_single, enumerate(cleaned))) - # Extract just the rewards for return - rewards = [r[0] for r in results] - - return rewards + return [r[0] for r in results] - -# ---- W&B init ---- +# Initialize W&B wandb.init( project=cfg.wandb_project, entity=cfg.wandb_entity, name=run_name, config={ - "model": MODEL_ID, + "model": cfg.model, "reward_config": reward_config.model_dump(), - "grpo_config": { + "training": { + "learning_rate": args.learning_rate, + "batch_size": args.per_device_train_batch_size, + "num_epochs": args.num_train_epochs, + }, + "grpo": { "beta": args.beta, "epsilon": args.epsilon, "temperature": args.temperature, @@ -190,17 +183,18 @@ def batch_reward_fn(prompts: List[str], completions: List[str], **kwargs) -> Lis }, ) -# ---- trainer (NO ref_model kwarg) ---- +# Initialize trainer trainer = GRPOTrainer( - model=MODEL_ID, + model=cfg.model, reward_funcs=[batch_reward_fn], args=args, train_dataset=train_ds, - eval_dataset=eval_ds if use_eval else None, + eval_dataset=eval_ds, processing_class=tok, callbacks=[reward_logger], ) +# Train and save trainer.train() trainer.save_model(args.output_dir) tok.save_pretrained(args.output_dir) diff --git a/src/runners/grpo_sweep.py b/src/runners/grpo_sweep.py new file mode 100644 index 0000000..845aad0 --- /dev/null +++ b/src/runners/grpo_sweep.py @@ -0,0 +1,191 @@ +""" +W&B Sweep runner for GRPO training. +This script is called by wandb agent with sweep parameters injected via wandb.config. +""" + +from datasets import load_dataset +from transformers import AutoTokenizer +from trl import GRPOTrainer, GRPOConfig +import torch +from src.config import Config +from src.rewards.bioinformatics.scorer import Scorer +from src.rewards.bioinformatics.reward_config import RewardConfig +from src.rewards.bioinformatics.logger import RewardComponentLogger +import datetime +from typing import List +import wandb +from concurrent.futures import ThreadPoolExecutor +from threading import Lock +import re + + +def load_train_val_datasets(cfg): + """Load and preprocess training and validation datasets.""" + def select_prompt_column(ds): + cols = set(ds.column_names) + keep_cols = ["prompt"] + [c for c in ["data_source", "ability", "reward_model", "extra_info"] if c in cols] + return ds.select_columns(keep_cols) + + train_ds = load_dataset("parquet", data_files=cfg.train_dataset, split="train") + val_ds = load_dataset("parquet", data_files=cfg.val_dataset, split="train") + + return select_prompt_column(train_ds), select_prompt_column(val_ds) + + +def main(): + """Run GRPO training with W&B sweep parameters.""" + # Initialize W&B run (sweep agent injects config) + run = wandb.init() + + # Get base config and sweep params + cfg = Config() + sweep_config = wandb.config + + run_name = f"sweep-{run.id}-{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" + + # Dataset loading + train_ds, eval_ds = load_train_val_datasets(cfg) + + # Tokenizer setup + tok = AutoTokenizer.from_pretrained(cfg.model, use_fast=True, trust_remote_code=True) + tok.padding_side = "left" + tok.eos_token = "" + tok.bos_token = "" + tok.pad_token = "[PAD]" + + assert tok.eos_token_id == 30001, f"Expected eos_token_id=30001, got {tok.eos_token_id}" + assert tok.bos_token_id == 30000, f"Expected bos_token_id=30000, got {tok.bos_token_id}" + assert tok.pad_token_id == 3, f"Expected pad_token_id=3, got {tok.pad_token_id}" + + model_init_kwargs = { + "trust_remote_code": True, + "eos_token_id": tok.eos_token_id, + "bos_token_id": tok.bos_token_id, + "pad_token_id": tok.pad_token_id, + } + + # Training configuration with sweep parameters + args = GRPOConfig( + model_init_kwargs=model_init_kwargs, + output_dir=f"/s3/checkpoints/verl-grpo-sweeps/{run_name}", + + # Training parameters (use sweep config or defaults) + num_train_epochs=1, # Short for sweeps + learning_rate=sweep_config.get("learning_rate", 3e-6), + lr_scheduler_type="constant", + warmup_ratio=0.0, + per_device_train_batch_size=sweep_config.get("per_device_train_batch_size", 16), + gradient_accumulation_steps=1, + max_steps=100, # Quick evaluation for sweeps + max_grad_norm=0.5, + seed=42, + + # Logging and checkpointing + save_strategy="no", # Don't save checkpoints for sweeps + logging_strategy="steps", + logging_steps=5, + report_to=["wandb"], + + # Evaluation + do_eval=True, + eval_strategy="steps", + eval_steps=100, + + # Optimization + bf16=torch.cuda.is_available(), + gradient_checkpointing=False, + + # GRPO-specific (use sweep config or defaults) + beta=sweep_config.get("beta", 1e-3), + epsilon=sweep_config.get("epsilon", 0.2), + loss_type="bnpo", + scale_rewards=True, + mask_truncated_completions=False, + disable_dropout=True, + + # Generation parameters (use sweep config or defaults) + remove_unused_columns=False, + max_prompt_length=1024, + num_generations=sweep_config.get("num_generations", 8), + max_completion_length=256, + temperature=sweep_config.get("temperature", 0.95), + top_p=sweep_config.get("top_p", 0.90), + + # vLLM configuration + use_vllm=True, + vllm_gpu_memory_utilization=0.15, + vllm_mode="colocate", + ) + + # Reward configuration with sweep parameters + reward_config = RewardConfig( + punish_mode=sweep_config.get("reward_punish_mode", False), + length_penalty=sweep_config.get("reward_length_penalty", True), + min_length=sweep_config.get("reward_min_length", 1000), + max_length=sweep_config.get("reward_max_length", 30000), + ori_min=1, + ori_max=1, + ori_weight=sweep_config.get("reward_ori_weight", 1.0), + promoter_min=1, + promoter_max=sweep_config.get("reward_promoter_max", 5), + promoter_weight=sweep_config.get("reward_promoter_weight", 1.0), + terminator_min=0, + terminator_max=sweep_config.get("reward_terminator_max", 2), + terminator_weight=sweep_config.get("reward_terminator_weight", 0.5), + marker_min=1, + marker_max=sweep_config.get("reward_marker_max", 2), + marker_weight=sweep_config.get("reward_marker_weight", 1.0), + cds_min=1, + cds_max=sweep_config.get("reward_cds_max", 5), + cds_weight=sweep_config.get("reward_cds_weight", 1.0), + location_aware=sweep_config.get("reward_location_aware", True), + ) + + # Initialize scorer and logger + scorer = Scorer(reward_config) + reward_logger = RewardComponentLogger(log_frequency=5) # Log frequently for short runs + component_lock = Lock() + + # Reward function + def score_single(idx_and_seq): + idx, seq = idx_and_seq + try: + score, components = scorer.score(seq) + with component_lock: + reward_logger.add_components(components, float(score)) + return float(score), components + except Exception as e: + print(f"Warning: Failed to score completion {idx} (len={len(seq)}): {str(e)[:100]}") + return 0.0, None + + def batch_reward_fn(prompts: List[str], completions: List[str], **kwargs) -> List[float]: + cleaned = [re.sub(r'[^ATCG]', '', c.upper().replace(" ", "")) for c in completions] + with ThreadPoolExecutor(max_workers=8) as executor: + results = list(executor.map(score_single, enumerate(cleaned))) + return [r[0] for r in results] + + # Initialize trainer + trainer = GRPOTrainer( + model=cfg.model, + reward_funcs=[batch_reward_fn], + args=args, + train_dataset=train_ds, + eval_dataset=eval_ds, + processing_class=tok, + callbacks=[reward_logger], + ) + + # Train + trainer.train() + + # Save best checkpoint info to W&B + trainer.save_model(args.output_dir) + tok.save_pretrained(args.output_dir) + + # Finish run + wandb.finish() + + +if __name__ == "__main__": + main() + diff --git a/sweep_config.yaml b/sweep_config.yaml new file mode 100644 index 0000000..188fed2 --- /dev/null +++ b/sweep_config.yaml @@ -0,0 +1,107 @@ +# W&B Sweep Configuration for PlasmidRL GRPO Training +# +# Usage: +# 1. Initialize the sweep: +# wandb sweep sweep_config.yaml +# +# 2. Copy the sweep ID from output (e.g., mcclain/plasmidrl-grpo-sweeps/abc123xyz) +# +# 3. Run agent(s) with Docker Compose: +# SWEEP_ID=mcclain/plasmidrl-grpo-sweeps/abc123xyz docker compose up grpo-sweep +# +# 4. Optional: Run multiple agents in parallel for faster sweeps: +# SWEEP_ID= docker compose up --scale grpo-sweep=3 +# +# Monitor progress at: https://wandb.ai/mcclain/plasmidrl-grpo-sweeps + +program: src/runners/grpo_sweep.py +method: bayes # Options: grid, random, bayes +metric: + name: reward_components/total_reward/mean + goal: maximize + +parameters: + # Training hyperparameters + learning_rate: + distribution: log_uniform_values + min: 1e-6 + max: 1e-4 + + per_device_train_batch_size: + values: [8, 16, 32] + + num_generations: + values: [4, 8, 16] + + temperature: + distribution: uniform + min: 0.7 + max: 1.4 + + top_p: + distribution: uniform + min: 0.85 + max: 0.95 + + # GRPO-specific parameters + beta: + distribution: log_uniform_values + min: 1e-4 + max: 1e-2 + + epsilon: + distribution: uniform + min: 0.1 + max: 0.3 + + # Reward config parameters + reward_punish_mode: + values: [true, false] + + reward_length_penalty: + values: [true, false] + + reward_min_length: + values: [500, 1000, 2000] + + reward_max_length: + values: [20000, 30000, 50000] + + reward_promoter_max: + values: [3, 5, 10] + + reward_terminator_max: + values: [1, 2, 3] + + reward_marker_max: + values: [1, 2, 3] + + reward_cds_max: + values: [3, 5, 10] + + reward_location_aware: + values: [true, false] + + # Reward component weights (0 = disabled) + reward_ori_weight: + values: [0.0, 0.5, 1.0, 2.0] + + reward_promoter_weight: + values: [0.0, 0.5, 1.0, 2.0] + + reward_terminator_weight: + values: [0.0, 0.25, 0.5, 1.0] + + reward_marker_weight: + values: [0.0, 0.5, 1.0, 2.0] + + reward_cds_weight: + values: [0.0, 0.5, 1.0, 2.0] + +# Early termination to stop poorly performing runs +early_terminate: + type: hyperband + min_iter: 20 # Minimum steps before termination + eta: 3 + s: 2 + From c1fc8056b712cfdda35081341ae478b9a2487d2f Mon Sep 17 00:00:00 2001 From: mcclain Date: Mon, 3 Nov 2025 11:46:13 +0000 Subject: [PATCH 09/11] slight reorg for sweeps --- docker-compose.yaml | 4 +- src/runners/grpo_sweep.py | 2 +- sweeps/README.md | 77 +++++++++++ SWEEPS.md => sweeps/SWEEPS.md | 63 ++++++++- .../configs/sweep_config.yaml | 4 +- sweeps/configs/sweep_config_refined.yaml | 128 ++++++++++++++++++ sweeps/configs/sweep_config_training.yaml | 112 +++++++++++++++ sweeps/run_sweep_agent.py | 12 ++ 8 files changed, 391 insertions(+), 11 deletions(-) create mode 100644 sweeps/README.md rename SWEEPS.md => sweeps/SWEEPS.md (54%) rename sweep_config.yaml => sweeps/configs/sweep_config.yaml (96%) create mode 100644 sweeps/configs/sweep_config_refined.yaml create mode 100644 sweeps/configs/sweep_config_training.yaml create mode 100755 sweeps/run_sweep_agent.py diff --git a/docker-compose.yaml b/docker-compose.yaml index edc2a03..be91ca5 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -71,14 +71,16 @@ services: command: > bash -lc " mkdir -p /tmp/wandb && \ + export PYTHONPATH=/mcclain && \ uv sync --frozen && \ - wandb agent ${SWEEP_ID} + uv run wandb agent ${SWEEP_ID} " environment: - WANDB_ENTITY=mcclain - WANDB_PROJECT=plasmidrl-grpo-sweeps - WANDB_DIR=/tmp/wandb - SWEEP_ID=${SWEEP_ID} + - PYTHONPATH=/mcclain # VERL: PPO trainer verl-ppo: diff --git a/src/runners/grpo_sweep.py b/src/runners/grpo_sweep.py index 845aad0..ef0106b 100644 --- a/src/runners/grpo_sweep.py +++ b/src/runners/grpo_sweep.py @@ -76,7 +76,7 @@ def main(): warmup_ratio=0.0, per_device_train_batch_size=sweep_config.get("per_device_train_batch_size", 16), gradient_accumulation_steps=1, - max_steps=100, # Quick evaluation for sweeps + max_steps=sweep_config.get("max_steps", 100), max_grad_norm=0.5, seed=42, diff --git a/sweeps/README.md b/sweeps/README.md new file mode 100644 index 0000000..b33321e --- /dev/null +++ b/sweeps/README.md @@ -0,0 +1,77 @@ +# PlasmidRL Hyperparameter Sweeps + +This directory contains all files related to W&B hyperparameter sweeps for GRPO training. + +## Structure + +``` +sweeps/ +├── configs/ # Sweep configuration files +│ ├── sweep_config.yaml # Full sweep (training + reward) +│ ├── sweep_config_training.yaml # Training hyperparameters only +│ └── sweep_config_refined.yaml # Refined ranges (500 steps) - RECOMMENDED +├── run_sweep_agent.py # Wrapper script for W&B agent +├── SWEEPS.md # Detailed documentation +└── README.md # This file +``` + +## Quick Start + +```bash +# 1. Initialize sweep (from project root) +wandb sweep sweeps/configs/sweep_config_refined.yaml + +# 2. Run agent with the returned SWEEP_ID +SWEEP_ID= docker compose up grpo-sweep + +# 3. Monitor at https://wandb.ai/mcclain/plasmidrl-grpo-sweeps +``` + +## Which Config to Use? + +- **`sweep_config_refined.yaml`** ⭐ - Use this after initial exploration to fine-tune + - 500 steps per trial for stable results + - Narrow ranges around best performers + - Fixed: batch_size=16, num_generations=8 + +- **`sweep_config_training.yaml`** - Use for broad training hyperparameter exploration + - 100 steps per trial for quick iteration + - Wide ranges for all training params + - Fixed reward configs + +- **`sweep_config.yaml`** - Use to explore reward configurations too + - 100 steps per trial + - Includes reward weight tuning + - Most comprehensive but slowest convergence + +## Related Files + +- **Sweep runner:** `src/runners/grpo_sweep.py` - Main training script called by W&B agent +- **Reward logger:** `src/rewards/bioinformatics/logger.py` - Logs detailed metrics to W&B +- **Docker service:** `docker-compose.yaml` (grpo-sweep service) + +## Documentation + +See [SWEEPS.md](./SWEEPS.md) for detailed documentation on: +- Configuration options +- Optimization strategy +- Best practices +- Troubleshooting + +## Sweep Results + +After running sweeps, document key findings here: + +### Sweep 1: Initial Exploration (Complete) +- **Best config:** Batch 16 + 8 generations +- **Learning rate:** 4e-05 to 8e-05 range optimal +- **Epsilon:** 0.27-0.30 (high exploration helps) +- **Beta:** 0.0008-0.0023 range works well +- **Success rate:** 84% (127/151 runs) +- **Best reward:** 0.86 + +### Sweep 2: Refined Training (In Progress) +- **Status:** Running +- **Config:** `sweep_config_refined.yaml` +- **Notes:** 500 steps per trial for more stable evaluation + diff --git a/SWEEPS.md b/sweeps/SWEEPS.md similarity index 54% rename from SWEEPS.md rename to sweeps/SWEEPS.md index cfaf340..88d393a 100644 --- a/SWEEPS.md +++ b/sweeps/SWEEPS.md @@ -7,12 +7,15 @@ This directory contains configuration for running hyperparameter sweeps using We ### 1. Initialize a Sweep ```bash -wandb sweep sweep_config.yaml +# From project root +wandb sweep sweeps/configs/sweep_config_refined.yaml ``` This will output a sweep ID like: `mcclain/plasmidrl-grpo-sweeps/abc123xyz` -### 2. Run Sweep Agent(s) +**Copy this sweep ID** - you'll need it for the next step. + +### 2. Run Sweep Agent(s) in Docker **Single agent:** ```bash @@ -29,9 +32,33 @@ SWEEP_ID=mcclain/plasmidrl-grpo-sweeps/abc123xyz docker compose up --scale grpo- Visit your W&B dashboard: - https://wandb.ai/mcclain/plasmidrl-grpo-sweeps -## Sweep Configuration +## Available Sweep Configurations + +All sweep configs are in `configs/`: + +### 1. `sweep_config_refined.yaml` (Recommended) +**Best for:** Finding optimal hyperparameters based on initial sweep insights +- **Duration:** 500 steps per trial (stable evaluation) +- **Focus:** Refined ranges around best performers +- **Fixed:** Batch size 16, 8 generations (best combo) +- **Tuned:** Learning rate, beta, epsilon, temperature, top_p + +### 2. `sweep_config_training.yaml` +**Best for:** Broad exploration of training hyperparameters +- **Duration:** 100 steps per trial (quick evaluation) +- **Focus:** Training parameters only +- **Fixed:** All reward configurations +- **Tuned:** Full ranges for LR, batch size, generations, etc. + +### 3. `sweep_config.yaml` +**Best for:** Full exploration including reward weights +- **Duration:** 100 steps per trial +- **Focus:** Both training and reward parameters +- **Tuned:** Everything (training + reward configs) + +## Sweep Configuration Details -The sweep is configured in `sweep_config.yaml` with the following parameters: +The sweep configurations include the following parameters: ### Training Hyperparameters - `learning_rate`: 1e-6 to 1e-4 (log uniform) @@ -63,9 +90,21 @@ The sweep uses **Bayesian optimization** to intelligently explore the hyperparam - **Early termination**: Hyperband stops poorly performing runs after 20 steps - **Smart search**: Focuses on promising hyperparameter regions +## Directory Structure + +``` +sweeps/ +├── configs/ # Sweep configuration files +│ ├── sweep_config.yaml # Full sweep (training + reward) +│ ├── sweep_config_training.yaml # Training hyperparameters only +│ └── sweep_config_refined.yaml # Refined ranges (500 steps) +├── run_sweep_agent.py # Wrapper script for W&B agent +└── SWEEPS.md # This file +``` + ## Customizing Sweeps -Edit `sweep_config.yaml` to: +Edit any config in `configs/` to: 1. Change the search method (`bayes`, `grid`, `random`) 2. Add/remove parameters 3. Adjust parameter ranges @@ -76,21 +115,31 @@ Edit `sweep_config.yaml` to: 1. **Start small**: Run a few trials manually to validate config 2. **Use multiple agents**: Parallel agents speed up sweeps significantly + - Each run takes ~5-10 minutes (100 steps) + - 3 parallel agents = ~50 trials in 3-4 hours 3. **Monitor early**: Check results after 5-10 runs to ensure valid config 4. **Stop bad sweeps**: If all runs fail, stop and fix the config 5. **Save good configs**: Note the best hyperparameters for future runs +6. **Full training**: Once you find best params, run full training (5000+ steps) with those settings ## Troubleshooting +**Sweep initialization fails:** +- Check W&B login: `wandb login` +- Verify you're logged into the correct W&B entity + **Sweep not starting:** -- Verify W&B login: `wandb login` -- Check SWEEP_ID format includes entity/project/id +- Verify SWEEP_ID format includes entity/project/id (e.g., `mcclain/plasmidrl-grpo-sweeps/abc123`) +- Check you copied the full sweep ID from initialization output +- Ensure `.env` file has `WANDB_API_KEY` **All runs failing:** - Check logs: `docker compose logs grpo-sweep` - Verify reward config constraints are achievable +- Run a single step manually: `docker compose run --rm grpo-sweep` **Out of memory:** - Reduce `per_device_train_batch_size` in sweep config - Reduce `num_generations` parameter range +- Stop other running containers first diff --git a/sweep_config.yaml b/sweeps/configs/sweep_config.yaml similarity index 96% rename from sweep_config.yaml rename to sweeps/configs/sweep_config.yaml index 188fed2..60ebf2d 100644 --- a/sweep_config.yaml +++ b/sweeps/configs/sweep_config.yaml @@ -6,7 +6,7 @@ # # 2. Copy the sweep ID from output (e.g., mcclain/plasmidrl-grpo-sweeps/abc123xyz) # -# 3. Run agent(s) with Docker Compose: +# 3. Run agent(s): # SWEEP_ID=mcclain/plasmidrl-grpo-sweeps/abc123xyz docker compose up grpo-sweep # # 4. Optional: Run multiple agents in parallel for faster sweeps: @@ -14,7 +14,7 @@ # # Monitor progress at: https://wandb.ai/mcclain/plasmidrl-grpo-sweeps -program: src/runners/grpo_sweep.py +program: sweeps/run_sweep_agent.py method: bayes # Options: grid, random, bayes metric: name: reward_components/total_reward/mean diff --git a/sweeps/configs/sweep_config_refined.yaml b/sweeps/configs/sweep_config_refined.yaml new file mode 100644 index 0000000..211cd7b --- /dev/null +++ b/sweeps/configs/sweep_config_refined.yaml @@ -0,0 +1,128 @@ +# W&B Sweep Configuration for PlasmidRL GRPO - Refined Training (500 steps) +# Focus: Refined hyperparameter search based on initial sweep findings +# Duration: 500 steps per trial for more stable evaluation +# +# Key Insights from Initial Sweep: +# - Batch size 16 + 8 generations is best (dominates top 15 runs) +# - Learning rate sweet spot: ~4e-05 to 8e-05 +# - High epsilon helps: 0.27-0.30 +# - Beta: 0.0008-0.0023 range appears frequently +# - Temperature: 1.07-1.39 all viable +# - Top-p: 0.86-0.93 clustering +# +# Usage: +# 1. Initialize the sweep: +# wandb sweep sweep_config_refined.yaml +# +# 2. Copy the sweep ID from output (e.g., mcclain/plasmidrl-grpo-sweeps/abc123xyz) +# +# 3. Run agent(s): +# SWEEP_ID=mcclain/plasmidrl-grpo-sweeps/abc123xyz docker compose up grpo-sweep +# +# Monitor progress at: https://wandb.ai/mcclain/plasmidrl-grpo-sweeps + +program: sweeps/run_sweep_agent.py +method: bayes +metric: + name: reward_components/total_reward/mean + goal: maximize + +parameters: + # ==================== TRAINING HYPERPARAMETERS (REFINED RANGES) ==================== + + max_steps: + value: 500 # Longer runs for stable evaluation + + # Learning rate - narrow range around sweet spot + learning_rate: + distribution: log_uniform_values + min: 3.5e-05 + max: 9.0e-05 + + # Batch size - focus on best performing config + per_device_train_batch_size: + values: [16] # 16 dominated top runs + + # Generations - focus on best performing config + num_generations: + values: [8] # 8 generations worked best + + # Temperature - narrow range around viable values + temperature: + distribution: uniform + min: 1.05 + max: 1.40 + + # Top-p - narrow range around cluster + top_p: + distribution: uniform + min: 0.86 + max: 0.93 + + # Beta - focus on frequently appearing range + beta: + distribution: log_uniform_values + min: 0.0008 + max: 0.0023 + + # Epsilon - favor high exploration + epsilon: + distribution: uniform + min: 0.27 + max: 0.30 + + # ==================== REWARD CONFIG (FIXED - BEST FROM INITIAL SWEEP) ==================== + + reward_punish_mode: + value: true + + reward_length_penalty: + value: false + + reward_min_length: + value: 1000 + + reward_max_length: + value: 30000 + + reward_promoter_max: + value: 5 + + reward_terminator_max: + value: 2 + + reward_marker_max: + value: 2 + + reward_cds_max: + value: 5 + + reward_location_aware: + value: true + + # Reward component weights (all enabled with standard weights) + reward_ori_weight: + value: 1.0 + + reward_promoter_weight: + value: 1.0 + + reward_terminator_weight: + value: 0.5 + + reward_marker_weight: + value: 1.0 + + reward_cds_weight: + value: 1.0 + +# Early termination - adjusted for longer runs +early_terminate: + type: hyperband + min_iter: 100 # Higher threshold for 500-step runs + eta: 3 + s: 2 + +# Optional: Limit total number of runs +# run_cap: 30 + diff --git a/sweeps/configs/sweep_config_training.yaml b/sweeps/configs/sweep_config_training.yaml new file mode 100644 index 0000000..f53aad7 --- /dev/null +++ b/sweeps/configs/sweep_config_training.yaml @@ -0,0 +1,112 @@ +# W&B Sweep Configuration for PlasmidRL GRPO Training +# Focus: Training hyperparameters only (reward config held constant) +# +# Usage: +# 1. Initialize the sweep: +# wandb sweep sweep_config_training.yaml +# +# 2. Copy the sweep ID from output (e.g., mcclain/plasmidrl-grpo-sweeps/abc123xyz) +# +# 3. Run agent(s): +# SWEEP_ID=mcclain/plasmidrl-grpo-sweeps/abc123xyz docker compose up grpo-sweep +# +# Monitor progress at: https://wandb.ai/mcclain/plasmidrl-grpo-sweeps + +program: sweeps/run_sweep_agent.py +method: bayes +metric: + name: reward_components/total_reward/mean + goal: maximize + +parameters: + # ==================== TRAINING HYPERPARAMETERS (SWEPT) ==================== + learning_rate: + distribution: log_uniform_values + min: 1e-6 + max: 1e-4 + + per_device_train_batch_size: + values: [4, 8, 16] + + num_generations: + values: [2, 4, 8] + + temperature: + distribution: uniform + min: 0.7 + max: 1.4 + + top_p: + distribution: uniform + min: 0.85 + max: 0.95 + + # GRPO-specific parameters + beta: + distribution: log_uniform_values + min: 1e-4 + max: 1e-2 + + epsilon: + distribution: uniform + min: 0.1 + max: 0.3 + + # ==================== REWARD CONFIG (FIXED) ==================== + # All reward parameters held constant for this sweep + + reward_punish_mode: + value: true + + reward_length_penalty: + value: false + + reward_min_length: + value: 1000 + + reward_max_length: + value: 30000 + + reward_promoter_max: + value: 5 + + reward_terminator_max: + value: 2 + + reward_marker_max: + value: 2 + + reward_cds_max: + value: 5 + + reward_location_aware: + value: true + + # Reward component weights (all enabled with standard weights) + reward_ori_weight: + value: 1.0 + + reward_promoter_weight: + value: 1.0 + + reward_terminator_weight: + value: 0.5 + + reward_marker_weight: + value: 1.0 + + reward_cds_weight: + value: 1.0 + +# Early termination to stop poorly performing runs +early_terminate: + type: hyperband + min_iter: 20 + eta: 3 + s: 2 + +# Optional: Limit total number of runs +# run_cap: 50 + + + diff --git a/sweeps/run_sweep_agent.py b/sweeps/run_sweep_agent.py new file mode 100755 index 0000000..7687db0 --- /dev/null +++ b/sweeps/run_sweep_agent.py @@ -0,0 +1,12 @@ +#!/usr/bin/env python +"""Wrapper script for W&B sweep agent.""" + +if __name__ == "__main__": + from src.runners.grpo_sweep import main + main() + + + + + + From c6c25caba8c408f1088c1b498958d1e60750bbf7 Mon Sep 17 00:00:00 2001 From: mcclain Date: Tue, 4 Nov 2025 11:21:24 +0000 Subject: [PATCH 10/11] trying with length --- AGENTS.md | 74 ++++++++- src/rewards/bioinformatics/reward_config.py | 16 +- src/rewards/bioinformatics/scorer.py | 53 ++++++- src/runners/grpo_sweep.py | 25 ++-- sweeps/README.md | 42 ++++-- sweeps/SWEEPS.md | 73 +++++++-- .../configs/sweep_config_length_reward.yaml | 124 ++++++++++++++++ .../sweep_config_training_with_length.yaml | 140 ++++++++++++++++++ 8 files changed, 497 insertions(+), 50 deletions(-) create mode 100644 sweeps/configs/sweep_config_length_reward.yaml create mode 100644 sweeps/configs/sweep_config_training_with_length.yaml diff --git a/AGENTS.md b/AGENTS.md index 194538a..cabe706 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -117,27 +117,41 @@ PlasmidRL/ │ ├── config.py # Pydantic config (loads .env) │ ├── runners/ # Training/inference entry points │ │ ├── grpo.py # TRL GRPO trainer +│ │ ├── grpo_sweep.py # W&B sweep runner for GRPO │ │ ├── es.py # TRL Evolution Strategies trainer │ │ └── generate_samples.py # vLLM-based inference │ ├── rewards/ # Reward computation & scoring │ │ ├── rewards.py # Reward signal definitions │ │ ├── plasmid_informatics.py # Plasmid analysis (validity, GFP, etc.) -│ │ └── verl_reward.py # VERL-compatible reward wrapper +│ │ ├── verl_reward.py # VERL-compatible reward wrapper +│ │ └── bioinformatics/ # Bioinformatics-based reward system +│ │ ├── scorer.py # Core scoring logic +│ │ ├── reward_config.py # Reward configuration +│ │ └── logger.py # W&B logging callback │ ├── eval/ # Evaluation utilities │ └── utils/ # Helper functions (model utils, S3 access) +├── sweeps/ # Hyperparameter sweep configurations +│ ├── configs/ # W&B sweep YAML configs +│ │ ├── sweep_config.yaml # Full sweep (training + reward) +│ │ ├── sweep_config_training.yaml # Training hyperparameters only +│ │ └── sweep_config_refined.yaml # Refined ranges (500 steps) +│ ├── run_sweep_agent.py # W&B agent wrapper script +│ ├── README.md # Quick reference +│ └── SWEEPS.md # Detailed sweep documentation ├── config/ # YAML configs for VERL trainers │ ├── verl_ppo.yaml │ ├── verl_grpo.yaml │ ├── verl_naive_ppo.yaml │ └── naive_grpo.yaml ├── docker/ # Specialized Dockerfiles -│ ├── verl.Dockerfile # VERL training environment -│ └── verl-monkey.patch # VERL patches/customizations +│ ├── verl.Dockerfile # VERL training environment +│ └── verl-monkey.patch # VERL patches/customizations ├── Dockerfile # Main development/TRL container ├── docker-compose.yaml # Multi-container orchestration -├── pyproject.toml # UV/Python project metadata -├── uv.lock # Frozen dependency lock file -└── README.md # User-facing documentation +├── pyproject.toml # UV/Python project metadata +├── uv.lock # Frozen dependency lock file +├── AGENTS.md # This file - AI agent briefing +└── README.md # User-facing documentation ``` ### Code Style @@ -382,6 +396,54 @@ uv run python -m src.main --help 3. **Review logs** from docker-compose or runner stderr 4. **Run with smaller dataset** or fewer steps for quick iteration +### Running Hyperparameter Sweeps + +**Location**: `sweeps/` + +W&B sweeps enable automated hyperparameter optimization. The project includes three pre-configured sweep strategies: + +1. **Refined sweep** (`sweeps/configs/sweep_config_refined.yaml`) - **Recommended** + - 500 steps per trial for stable evaluation + - Narrow ranges around best performers from initial exploration + - Fixed batch_size=16, num_generations=8 (proven best combo) + +2. **Training-only sweep** (`sweeps/configs/sweep_config_training.yaml`) + - 100 steps per trial for quick iteration + - Explores all training hyperparameters + - Fixed reward configurations + +3. **Full sweep** (`sweeps/configs/sweep_config.yaml`) + - 100 steps per trial + - Includes both training and reward parameter tuning + - Most comprehensive but slower convergence + +**Workflow:** + +```bash +# 1. Initialize sweep (from project root) +wandb sweep sweeps/configs/sweep_config_refined.yaml + +# 2. Run agent(s) with returned SWEEP_ID +SWEEP_ID= docker compose up grpo-sweep + +# 3. Optional: Run multiple parallel agents for faster sweeps +SWEEP_ID= docker compose up --scale grpo-sweep=3 + +# 4. Monitor at https://wandb.ai/mcclain/plasmidrl-grpo-sweeps +``` + +**Key files:** +- `sweeps/configs/` - Sweep configuration YAMLs +- `sweeps/run_sweep_agent.py` - W&B agent wrapper +- `src/runners/grpo_sweep.py` - Main training script called by W&B +- `sweeps/SWEEPS.md` - Detailed sweep documentation + +**Tips:** +- Start with `sweep_config_refined.yaml` if you've done initial exploration +- Use multiple agents (`--scale grpo-sweep=3`) to parallelize trials +- Each 100-step trial takes ~5-10 minutes; 500-step trials take ~25-50 minutes +- Document findings in `sweeps/README.md` for future reference + --- ## References & Links diff --git a/src/rewards/bioinformatics/reward_config.py b/src/rewards/bioinformatics/reward_config.py index 0ed3630..5de9774 100644 --- a/src/rewards/bioinformatics/reward_config.py +++ b/src/rewards/bioinformatics/reward_config.py @@ -3,14 +3,20 @@ class RewardConfig(BaseModel): - punish_mode: bool = True # penaize violations of the reward config as opposed to just not rewarding them - length_penalty: bool = False # penalize sequences that are too long or too short - min_length: Optional[int] = None - max_length: Optional[int] = None - location_aware: bool = True # reward sequences that are located in the correct location (e.g. prompoter then cds then terminatr) + punish_mode: bool = True # penalize violations of the reward config as opposed to just not rewarding them + length_reward_mode: bool = False # reward sequences based on length (replaces length_penalty) + min_length: Optional[int] = None # minimum acceptable length + max_length: Optional[int] = None # maximum acceptable length + ideal_min_length: Optional[int] = None # ideal minimum length for bonus reward + ideal_max_length: Optional[int] = None # ideal maximum length for bonus reward + length_reward_bonus: float = 0.5 # bonus multiplier for sequences in ideal length range + location_aware: bool = True # reward sequences that are located in the correct location (e.g. promoter then cds then terminator) # Penalty factor applied when min/max constraints are violated (outside of range) violation_penalty_factor: float = 0.5 + # Deprecated - use length_reward_mode instead + length_penalty: bool = False + ori_min: int = 1 ori_max: int = 1 allowed_oris: Optional[List[str]] = None diff --git a/src/rewards/bioinformatics/scorer.py b/src/rewards/bioinformatics/scorer.py index cf8fd3e..eddf89a 100644 --- a/src/rewards/bioinformatics/scorer.py +++ b/src/rewards/bioinformatics/scorer.py @@ -303,18 +303,55 @@ def score_cds(self, seq: str, annotations: Any) -> float: return max(0.0, min(1.0, base + cassette_bonus)) def score_length(self, seq: str, annotations: Any) -> float: - """Score sequence length (penalty if outside allowed range).""" - if not self.reward_config.length_penalty: + """ + Score sequence length with rewards for being within target range. + + If length_reward_mode=True: + - Sequences within [min_length, max_length] get full credit (1.0) + - Sequences within [ideal_min, ideal_max] get a bonus (up to 1.0 + length_reward_bonus) + - Sequences outside range get penalty or 0 based on punish_mode + + If length_reward_mode=False: + - Returns 1.0 (no length-based reward/penalty) + """ + if not self.reward_config.length_reward_mode: return 1.0 + L = len(seq) mn = self.reward_config.min_length mx = self.reward_config.max_length - assert mn is not None or mx is not None, "min_length or max_length must be set if length_penalty is True" - # Inside range: full credit; outside: penalty if punish_mode - in_range = (mn is None or L >= mn) and (mx is None or L <= mx) - if in_range: - return 1.0 - return float(self.reward_config.violation_penalty_factor) if self.reward_config.punish_mode else 1.0 + assert mn is not None and mx is not None, "min_length and max_length must be set if length_reward_mode is True" + + # Outside acceptable range: penalty + if L < mn or L > mx: + return float(self.reward_config.violation_penalty_factor) if self.reward_config.punish_mode else 0.5 + + # Within acceptable range: base credit of 1.0 + base_score = 1.0 + + # Check for ideal range bonus + ideal_min = self.reward_config.ideal_min_length + ideal_max = self.reward_config.ideal_max_length + + if ideal_min is not None and ideal_max is not None: + # If within ideal range, give full bonus + if ideal_min <= L <= ideal_max: + return base_score + self.reward_config.length_reward_bonus + + # If between min-ideal_min or ideal_max-max, give partial bonus + # based on distance from ideal range + if L < ideal_min: + # Between min and ideal_min + distance_ratio = (L - mn) / (ideal_min - mn) if ideal_min > mn else 1.0 + bonus = distance_ratio * self.reward_config.length_reward_bonus + return base_score + bonus + else: + # Between ideal_max and max + distance_ratio = (mx - L) / (mx - ideal_max) if mx > ideal_max else 1.0 + bonus = distance_ratio * self.reward_config.length_reward_bonus + return base_score + bonus + + return base_score def score(self, seq: str, source: str | None = None) -> Tuple[float, Dict[str, float]]: """ diff --git a/src/runners/grpo_sweep.py b/src/runners/grpo_sweep.py index ef0106b..44d2ec3 100644 --- a/src/runners/grpo_sweep.py +++ b/src/runners/grpo_sweep.py @@ -34,14 +34,15 @@ def select_prompt_column(ds): def main(): """Run GRPO training with W&B sweep parameters.""" - # Initialize W&B run (sweep agent injects config) - run = wandb.init() - - # Get base config and sweep params + # Get base config first cfg = Config() - sweep_config = wandb.config - run_name = f"sweep-{run.id}-{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" + # Initialize W&B run with a meaningful name + run_name = f"grpo-sweep-{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" + run = wandb.init(name=run_name) + + # Get sweep params + sweep_config = wandb.config # Dataset loading train_ds, eval_ds = load_train_val_datasets(cfg) @@ -65,9 +66,10 @@ def main(): } # Training configuration with sweep parameters + checkpoint_dir = f"/mnt/s3/phd-research-storage-1758274488/checkpoints/grpo-sweeps/{run_name}" args = GRPOConfig( model_init_kwargs=model_init_kwargs, - output_dir=f"/s3/checkpoints/verl-grpo-sweeps/{run_name}", + output_dir=checkpoint_dir, # Training parameters (use sweep config or defaults) num_train_epochs=1, # Short for sweeps @@ -81,7 +83,9 @@ def main(): seed=42, # Logging and checkpointing - save_strategy="no", # Don't save checkpoints for sweeps + save_strategy="steps", + save_steps=100, # Save every 100 steps + save_total_limit=2, # Keep only the last 2 checkpoints logging_strategy="steps", logging_steps=5, report_to=["wandb"], @@ -120,9 +124,12 @@ def main(): # Reward configuration with sweep parameters reward_config = RewardConfig( punish_mode=sweep_config.get("reward_punish_mode", False), - length_penalty=sweep_config.get("reward_length_penalty", True), + length_reward_mode=sweep_config.get("reward_length_reward_mode", False), min_length=sweep_config.get("reward_min_length", 1000), max_length=sweep_config.get("reward_max_length", 30000), + ideal_min_length=sweep_config.get("reward_ideal_min_length", None), + ideal_max_length=sweep_config.get("reward_ideal_max_length", None), + length_reward_bonus=sweep_config.get("reward_length_reward_bonus", 0.5), ori_min=1, ori_max=1, ori_weight=sweep_config.get("reward_ori_weight", 1.0), diff --git a/sweeps/README.md b/sweeps/README.md index b33321e..2bd264f 100644 --- a/sweeps/README.md +++ b/sweeps/README.md @@ -6,40 +6,58 @@ This directory contains all files related to W&B hyperparameter sweeps for GRPO ``` sweeps/ -├── configs/ # Sweep configuration files -│ ├── sweep_config.yaml # Full sweep (training + reward) -│ ├── sweep_config_training.yaml # Training hyperparameters only -│ └── sweep_config_refined.yaml # Refined ranges (500 steps) - RECOMMENDED -├── run_sweep_agent.py # Wrapper script for W&B agent -├── SWEEPS.md # Detailed documentation -└── README.md # This file +├── configs/ # Sweep configuration files +│ ├── sweep_config.yaml # Full sweep (training + reward) +│ ├── sweep_config_training.yaml # Training hyperparameters only +│ ├── sweep_config_refined.yaml # Refined ranges (500 steps) +│ ├── sweep_config_training_with_length.yaml # Training + length rewards ⭐ NEW +│ └── sweep_config_length_reward.yaml # Length reward testing only +├── run_sweep_agent.py # Wrapper script for W&B agent +├── SWEEPS.md # Detailed documentation +└── README.md # This file ``` ## Quick Start ```bash # 1. Initialize sweep (from project root) -wandb sweep sweeps/configs/sweep_config_refined.yaml +wandb sweep sweeps/configs/sweep_config_training_with_length.yaml # 2. Run agent with the returned SWEEP_ID SWEEP_ID= docker compose up grpo-sweep +# Or run multiple agents in parallel for faster sweeps: +SWEEP_ID= docker compose up --scale grpo-sweep=3 + # 3. Monitor at https://wandb.ai/mcclain/plasmidrl-grpo-sweeps ``` ## Which Config to Use? -- **`sweep_config_refined.yaml`** ⭐ - Use this after initial exploration to fine-tune +- **`sweep_config_training_with_length.yaml`** ⭐ **RECOMMENDED FOR NEXT SWEEP** + - 500 steps per trial for stable evaluation + - Broad training hyperparameter search + - Tests 2 length reward configurations: + - Smaller plasmids: 2-15kb (ideal: 3-12kb) + - Larger plasmids: 5-30kb (ideal: 7-20kb) + - Great for finding optimal hyperparameters with length rewards + +- **`sweep_config_refined.yaml`** - Use for fine-tuning without length rewards - 500 steps per trial for stable results - - Narrow ranges around best performers + - Narrow ranges around best performers from initial sweep - Fixed: batch_size=16, num_generations=8 -- **`sweep_config_training.yaml`** - Use for broad training hyperparameter exploration +- **`sweep_config_length_reward.yaml`** - Use to explore length reward ranges only + - 500 steps per trial + - Fixed: best training hyperparameters + - Variable: ideal length ranges and bonus multipliers + +- **`sweep_config_training.yaml`** - Use for broad exploration (no length rewards) - 100 steps per trial for quick iteration - Wide ranges for all training params - Fixed reward configs -- **`sweep_config.yaml`** - Use to explore reward configurations too +- **`sweep_config.yaml`** - Use to explore reward weights - 100 steps per trial - Includes reward weight tuning - Most comprehensive but slowest convergence diff --git a/sweeps/SWEEPS.md b/sweeps/SWEEPS.md index 88d393a..c023949 100644 --- a/sweeps/SWEEPS.md +++ b/sweeps/SWEEPS.md @@ -43,14 +43,32 @@ All sweep configs are in `configs/`: - **Fixed:** Batch size 16, 8 generations (best combo) - **Tuned:** Learning rate, beta, epsilon, temperature, top_p -### 2. `sweep_config_training.yaml` +### 2. `sweep_config_training_with_length.yaml` (RECOMMENDED FOR NEXT SWEEP) +**Best for:** Full hyperparameter search with length rewards +- **Duration:** 500 steps per trial +- **Focus:** All training hyperparameters + length rewards +- **Fixed:** Standard reward component weights +- **Tuned:** LR, batch size, generations, temperature, top_p, beta, epsilon, length bonus +- **Length configs:** Tests 2 combinations: + - Smaller plasmids: min=2000, ideal=3000-12000, max=15000 + - Larger plasmids: min=5000, ideal=7000-20000, max=30000 + +### 3. `sweep_config_length_reward.yaml` +**Best for:** Testing length-based reward strategies only +- **Duration:** 500 steps per trial +- **Focus:** Length reward parameters (ideal ranges, bonus multipliers) +- **Fixed:** Best training hyperparameters from previous sweeps +- **Tuned:** Ideal length ranges, length reward bonus +- **Feature:** Rewards sequences within ideal length ranges with bonuses + +### 4. `sweep_config_training.yaml` **Best for:** Broad exploration of training hyperparameters - **Duration:** 100 steps per trial (quick evaluation) - **Focus:** Training parameters only - **Fixed:** All reward configurations - **Tuned:** Full ranges for LR, batch size, generations, etc. -### 3. `sweep_config.yaml` +### 5. `sweep_config.yaml` **Best for:** Full exploration including reward weights - **Duration:** 100 steps per trial - **Focus:** Both training and reward parameters @@ -79,8 +97,40 @@ The sweep configurations include the following parameters: - `reward_terminator_weight`: 0.0, 0.25, 0.5, or 1.0 - `reward_marker_weight`: 0.0, 0.5, 1.0, or 2.0 - `reward_cds_weight`: 0.0, 0.5, 1.0, or 2.0 -- **Penalties**: `reward_punish_mode`, `reward_length_penalty` -- **Features**: `reward_location_aware` (cassette bonuses) +- **Penalties**: `reward_punish_mode` +- **Features**: + - `reward_location_aware` (cassette bonuses) + - `reward_length_reward_mode` (length-based rewards) + +### Length Reward System (NEW) +When `reward_length_reward_mode` is enabled: +- **Base reward (1.0)**: Sequences within `[min_length, max_length]` +- **Bonus reward**: Sequences within `[ideal_min_length, ideal_max_length]` get up to `1.0 + length_reward_bonus` +- **Partial bonus**: Sequences between min/ideal or ideal/max get proportional bonus +- **Penalty**: Sequences outside acceptable range get `violation_penalty_factor` or 0.5 + +Example configuration: +```yaml +reward_length_reward_mode: true +reward_min_length: 1000 # Minimum acceptable +reward_max_length: 30000 # Maximum acceptable +reward_ideal_min_length: 3000 # Ideal minimum +reward_ideal_max_length: 10000 # Ideal maximum +reward_length_reward_bonus: 0.5 # +50% bonus for ideal range +``` + +## Checkpoint Management + +Each sweep run automatically saves checkpoints to S3: +- **Location**: `/mnt/s3/phd-research-storage-1758274488/checkpoints/grpo-sweeps/{run_name}/` +- **Naming**: Run name matches W&B run name (e.g., `grpo-sweep-20251103_153516`) +- **Strategy**: Saves every 100 steps, keeps last 2 checkpoints +- **Contents**: Model weights + tokenizer + +To find the best checkpoint from a sweep: +1. Identify the best run in W&B dashboard +2. Note the run name (visible in W&B UI) +3. Find checkpoint at `/mnt/s3/.../checkpoints/grpo-sweeps/{run_name}/` ## Optimization Strategy @@ -94,12 +144,15 @@ The sweep uses **Bayesian optimization** to intelligently explore the hyperparam ``` sweeps/ -├── configs/ # Sweep configuration files -│ ├── sweep_config.yaml # Full sweep (training + reward) -│ ├── sweep_config_training.yaml # Training hyperparameters only -│ └── sweep_config_refined.yaml # Refined ranges (500 steps) -├── run_sweep_agent.py # Wrapper script for W&B agent -└── SWEEPS.md # This file +├── configs/ # Sweep configuration files +│ ├── sweep_config.yaml # Full sweep (training + reward) +│ ├── sweep_config_training.yaml # Training hyperparameters only +│ ├── sweep_config_refined.yaml # Refined ranges (500 steps) +│ ├── sweep_config_training_with_length.yaml # Training + length rewards (RECOMMENDED) +│ └── sweep_config_length_reward.yaml # Length reward testing only +├── run_sweep_agent.py # Wrapper script for W&B agent +├── README.md # Quick reference +└── SWEEPS.md # This file - detailed documentation ``` ## Customizing Sweeps diff --git a/sweeps/configs/sweep_config_length_reward.yaml b/sweeps/configs/sweep_config_length_reward.yaml new file mode 100644 index 0000000..1ef0e3c --- /dev/null +++ b/sweeps/configs/sweep_config_length_reward.yaml @@ -0,0 +1,124 @@ +# W&B Sweep Configuration for PlasmidRL GRPO - Length Reward Testing +# Focus: Test length-based rewards with different ideal ranges +# Duration: 500 steps per trial for stable evaluation +# +# Usage: +# 1. Initialize the sweep: +# wandb sweep sweeps/configs/sweep_config_length_reward.yaml +# +# 2. Copy the sweep ID from output (e.g., mcclain/plasmidrl-grpo-sweeps/abc123xyz) +# +# 3. Run agent(s): +# SWEEP_ID=mcclain/plasmidrl-grpo-sweeps/abc123xyz docker compose up grpo-sweep +# +# Monitor progress at: https://wandb.ai/mcclain/plasmidrl-grpo-sweeps + +program: sweeps/run_sweep_agent.py +method: bayes +metric: + name: reward_components/total_reward/mean + goal: maximize + +parameters: + # ==================== TRAINING HYPERPARAMETERS (FIXED - BEST FROM PREVIOUS SWEEPS) ==================== + + max_steps: + value: 500 # Longer runs for stable evaluation + + # Best learning rate from previous sweep + learning_rate: + value: 5.68e-05 + + # Best batch size + generations combo + per_device_train_batch_size: + value: 16 + + num_generations: + value: 8 + + # Best temperature from previous sweep + temperature: + value: 1.07 + + # Best top-p from previous sweep + top_p: + value: 0.904 + + # Best beta from previous sweep + beta: + value: 0.0016 + + # Best epsilon from previous sweep + epsilon: + value: 0.295 + + # ==================== LENGTH REWARD PARAMETERS (VARIABLE) ==================== + + reward_length_reward_mode: + value: true # Enable length-based rewards + + reward_min_length: + value: 1000 # Minimum acceptable plasmid length + + reward_max_length: + value: 30000 # Maximum acceptable plasmid length + + # Ideal length range for bonus rewards - sweep over different ranges + reward_ideal_min_length: + values: [2000, 3000, 4000, 5000] # Test different ideal minimums + + reward_ideal_max_length: + values: [8000, 10000, 12000, 15000] # Test different ideal maximums + + # Bonus multiplier for being in ideal range + reward_length_reward_bonus: + distribution: uniform + min: 0.2 + max: 1.0 + + # ==================== REWARD CONFIG (FIXED - BEST FROM PREVIOUS SWEEPS) ==================== + + reward_punish_mode: + value: true + + reward_promoter_max: + value: 5 + + reward_terminator_max: + value: 2 + + reward_marker_max: + value: 2 + + reward_cds_max: + value: 5 + + reward_location_aware: + value: true + + # Reward component weights (all enabled with standard weights) + reward_ori_weight: + value: 1.0 + + reward_promoter_weight: + value: 1.0 + + reward_terminator_weight: + value: 0.5 + + reward_marker_weight: + value: 1.0 + + reward_cds_weight: + value: 1.0 + +# Early termination - adjusted for longer runs +early_terminate: + type: hyperband + min_iter: 100 # Higher threshold for 500-step runs + eta: 3 + s: 2 + +# Optional: Limit total number of runs +# run_cap: 30 + diff --git a/sweeps/configs/sweep_config_training_with_length.yaml b/sweeps/configs/sweep_config_training_with_length.yaml new file mode 100644 index 0000000..5006440 --- /dev/null +++ b/sweeps/configs/sweep_config_training_with_length.yaml @@ -0,0 +1,140 @@ +# W&B Sweep Configuration for PlasmidRL GRPO - Training + Length Reward +# Focus: Broad training hyperparameter search with length-based rewards +# Duration: 500 steps per trial for stable evaluation +# +# Strategy: Test 2 length configurations across full hyperparameter space +# +# Usage: +# 1. Initialize the sweep: +# wandb sweep sweeps/configs/sweep_config_training_with_length.yaml +# +# 2. Copy the sweep ID from output (e.g., mcclain/plasmidrl-grpo-sweeps/abc123xyz) +# +# 3. Run agent(s): +# SWEEP_ID=mcclain/plasmidrl-grpo-sweeps/abc123xyz docker compose up grpo-sweep +# +# Monitor progress at: https://wandb.ai/mcclain/plasmidrl-grpo-sweeps + +program: sweeps/run_sweep_agent.py +method: bayes +metric: + name: reward_components/total_reward/mean + goal: maximize + +parameters: + # ==================== TRAINING HYPERPARAMETERS (BROAD SEARCH) ==================== + + max_steps: + value: 500 # Longer runs for stable evaluation + + # Learning rate - broad exploration + learning_rate: + distribution: log_uniform_values + min: 1e-6 + max: 1e-4 + + # Batch size + per_device_train_batch_size: + values: [8, 16, 32] + + # Generations per batch + num_generations: + values: [4, 8, 16] + + # Temperature - sampling randomness + temperature: + distribution: uniform + min: 0.7 + max: 1.4 + + # Top-p - nucleus sampling + top_p: + distribution: uniform + min: 0.85 + max: 0.95 + + # Beta - KL penalty coefficient + beta: + distribution: log_uniform_values + min: 1e-4 + max: 1e-2 + + # Epsilon - PPO-style clipping + epsilon: + distribution: uniform + min: 0.1 + max: 0.3 + + # ==================== LENGTH REWARD PARAMETERS (2 COMBOS) ==================== + + reward_length_reward_mode: + value: true # Enable length-based rewards + + # Test 2 length configurations + # Combo 1: min=2000, ideal_min=3000, ideal_max=12000, max=15000 + # Combo 2: min=5000, ideal_min=7000, ideal_max=20000, max=30000 + + reward_min_length: + values: [2000, 5000] + + reward_max_length: + values: [15000, 30000] + + reward_ideal_min_length: + values: [3000, 7000] + + reward_ideal_max_length: + values: [12000, 20000] + + # Bonus multiplier for being in ideal range + reward_length_reward_bonus: + distribution: uniform + min: 0.3 + max: 0.8 + + # ==================== REWARD CONFIG (FIXED - STANDARD) ==================== + + reward_punish_mode: + value: true + + reward_promoter_max: + value: 5 + + reward_terminator_max: + value: 2 + + reward_marker_max: + value: 2 + + reward_cds_max: + value: 5 + + reward_location_aware: + value: true + + # Reward component weights (all enabled with standard weights) + reward_ori_weight: + value: 1.0 + + reward_promoter_weight: + value: 1.0 + + reward_terminator_weight: + value: 0.5 + + reward_marker_weight: + value: 1.0 + + reward_cds_weight: + value: 1.0 + +# Early termination - adjusted for longer runs +early_terminate: + type: hyperband + min_iter: 100 # Higher threshold for 500-step runs + eta: 3 + s: 2 + +# Optional: Limit total number of runs +# run_cap: 50 + From 5be00c8b15a41c520481acc9a67f86964124c1af Mon Sep 17 00:00:00 2001 From: mcclain Date: Fri, 7 Nov 2025 14:06:48 +0000 Subject: [PATCH 11/11] added sweep info and better logging --- docker-compose.yaml | 8 +- src/config.py | 10 + src/eval/eval.py | 365 ++++++++++++++++++ src/eval/eval_config.py | 33 ++ src/rewards/bioinformatics/reward_config.py | 4 +- src/runners/grpo.py | 125 ++++-- src/runners/grpo_sweep.py | 50 ++- src/utils/training_utils.py | 243 ++++++++++++ .../sweep_config_training_with_eval.yaml | 144 +++++++ 9 files changed, 946 insertions(+), 36 deletions(-) create mode 100644 src/eval/eval.py create mode 100644 src/eval/eval_config.py create mode 100644 src/utils/training_utils.py create mode 100644 sweeps/configs/sweep_config_training_with_eval.yaml diff --git a/docker-compose.yaml b/docker-compose.yaml index be91ca5..d8f1338 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -41,8 +41,8 @@ services: " environment: # W&B settings for this job - - WANDB_ENTITY=mcclain - - WANDB_PROJECT=plasmidrl-trl-grpo + - WANDB_ENTITY=ucl-cssb + - WANDB_PROJECT=PlasmidRL - WANDB_TAGS=["plasmid","rl","trl","grpo"] - WANDB_NOTES=TRL GRPO training on plasmid design - WANDB_DIR=/tmp/wandb @@ -76,8 +76,8 @@ services: uv run wandb agent ${SWEEP_ID} " environment: - - WANDB_ENTITY=mcclain - - WANDB_PROJECT=plasmidrl-grpo-sweeps + - WANDB_ENTITY=ucl-cssb + - WANDB_PROJECT=PlasmidRL - WANDB_DIR=/tmp/wandb - SWEEP_ID=${SWEEP_ID} - PYTHONPATH=/mcclain diff --git a/src/config.py b/src/config.py index 872d0b2..1b1ad0d 100644 --- a/src/config.py +++ b/src/config.py @@ -47,6 +47,16 @@ class Config(BaseSettings): region_name: str = "us-east-1" runs_path: str = "runs/" infered_path: str = "infered/" + checkpoints_path: str = "checkpoints/" # S3 prefix for checkpoint storage + + # Production GRPO hyperparameters (from sweep optimization) + grpo_learning_rate: float = 0.00001906419115928539 + grpo_per_device_train_batch_size: int = 16 + grpo_num_generations: int = 4 + grpo_temperature: float = 1.2292317925218237 + grpo_top_p: float = 0.9086524230707756 + grpo_beta: float = 0.00088482365318492 + grpo_epsilon: float = 0.2649093053949679 model_config = { "env_file": ".env", diff --git a/src/eval/eval.py b/src/eval/eval.py new file mode 100644 index 0000000..03639b3 --- /dev/null +++ b/src/eval/eval.py @@ -0,0 +1,365 @@ +from vllm import LLM, SamplingParams +from typing import Optional, List, Dict, Any +from src.eval.eval_config import EvalConfig +from src.utils.training_utils import EvalRunner +from src.config import Config +import pandas as pd +import os +import plasmidkit as pk +import re + + +class _Feat: + """Simple feature container for annotation merging (same as Scorer).""" + def __init__(self, type: str, id: str | None, start: int, end: int, strand: str | None, evidence: Any = None): + self.type = type + self.id = id + self.start = int(start) + self.end = int(end) + self.strand = strand or "+" + self.evidence = evidence or {} + + +class SequenceAnalyzer: + """ + Analyzes plasmid sequences and extracts detailed annotation information. + + Similar to Scorer but returns detailed information instead of scores. + Uses plasmidkit for annotation and extracts counts and IDs for each feature type. + """ + + def __init__(self, eval_config: EvalConfig): + self.eval_config = eval_config + + def annotate(self, sequence: str) -> List[Any]: + """Annotate sequence with plasmidkit and merge overlapping features.""" + assert sequence, "sequence cannot be empty" + raw = pk.annotate(sequence, is_sequence=True) + return self._preprocess_annotations(raw) + + @staticmethod + def _overlap_len(a: Any, b: Any) -> int: + """Calculate overlap length between two features.""" + s1, e1 = int(a.start), int(a.end) + s2, e2 = int(b.start), int(b.end) + lo = max(min(s1, e1), min(s2, e2)) + hi = min(max(s1, e1), max(s2, e2)) + return max(0, hi - lo) + + def _to_feat(self, x: Any) -> _Feat: + """Convert annotation object to internal _Feat representation.""" + return _Feat( + type=x.type.lower() if x.type else "", + id=x.id if hasattr(x, "id") else None, + start=int(x.start), + end=int(x.end), + strand=x.strand if hasattr(x, "strand") else "+", + evidence=x.evidence if hasattr(x, "evidence") else {}, + ) + + def _merge_group(self, feats: List[Any], threshold: float, *, respect_strand: bool) -> List[_Feat]: + """Merge overlapping features of the same type based on overlap threshold.""" + if not feats: + return [] + items = [self._to_feat(f) for f in feats] + items.sort(key=lambda f: (f.strand, f.start, f.end)) + merged: List[_Feat] = [] + cur = items[0] + for nxt in items[1:]: + ovl = self._overlap_len(cur, nxt) + cur_len = max(0, cur.end - cur.start) + nxt_len = max(0, nxt.end - nxt.start) + min_len = max(1, min(cur_len, nxt_len)) + strands_compatible = (cur.strand == nxt.strand) or (not respect_strand) + if ovl / float(min_len) >= threshold and strands_compatible: + cur.start = min(cur.start, nxt.start) + cur.end = max(cur.end, nxt.end) + cur.id = f"{cur.id}|{nxt.id}" if cur.id or nxt.id else None + else: + merged.append(cur) + cur = nxt + merged.append(cur) + return merged + + def _preprocess_annotations(self, annotations: Any) -> List[Any]: + """ + Merge overlapping annotations and filter out CDS overlapping with other feature types. + + Same logic as Scorer._preprocess_annotations. + """ + feats = list(annotations) + thr = float(self.eval_config.overlap_merge_threshold) + type_key = lambda x: x.type.lower() if x.type else "" + + # Collect groups by type + groups: Dict[str, List[Any]] = {} + for f in feats: + groups.setdefault(type_key(f), []).append(f) + + # Merge per group for relevant types + merged_groups: Dict[str, List[_Feat]] = {} + for t in ("rep_origin", "ori", "origin_of_replication", "promoter", "terminator", "marker", "cds"): + if t in groups: + respect = t not in ("rep_origin", "ori", "origin_of_replication", "marker") + merged_groups[t] = self._merge_group(groups[t], thr, respect_strand=respect) + + # Suppress CDS if overlaps any non-CDS + non_cds: List[_Feat] = [] + for t in ("rep_origin", "ori", "origin_of_replication", "promoter", "terminator", "marker"): + non_cds.extend(merged_groups.get(t, [])) + + filtered_cds: List[_Feat] = [] + for c in merged_groups.get("cds", []): + if any(self._overlap_len(c, o) > 0 for o in non_cds): + continue + filtered_cds.append(c) + merged_groups["cds"] = filtered_cds + + # Rebuild final list + final: List[Any] = [] + merged_types = set(merged_groups.keys()) + for t, items in merged_groups.items(): + final.extend(items) + for f in feats: + t = type_key(f) + if t not in merged_types: + final.append(f) + return final + + def analyze(self, sequence: str) -> Dict[str, Any]: + """ + Analyze a sequence and extract detailed annotation information. + + Args: + sequence: DNA sequence to analyze + + Returns: + Dictionary with counts and IDs for each feature type + """ + annotations = self.annotate(sequence) + feats = list(annotations) + type_key = lambda x: x.type.lower() if x.type else "" + + # Extract features by type (no filtering - report all IDs found) + oris = [x for x in feats if type_key(x) in ("rep_origin", "ori", "origin_of_replication")] + promoters = [x for x in feats if type_key(x) == "promoter"] + terminators = [x for x in feats if type_key(x) == "terminator"] + markers = [x for x in feats if type_key(x) == "marker"] + cdss = [x for x in feats if type_key(x) == "cds"] + + # Extract IDs (handle merged IDs separated by |) + def extract_ids(features: List[Any]) -> List[str]: + ids = [] + for f in features: + if hasattr(f, "id") and f.id: + # Split merged IDs (separated by |) + ids.extend([id.strip() for id in str(f.id).split("|") if id.strip()]) + return ids + + return { + "ori_count": len(oris), + "ori_ids": ",".join(extract_ids(oris)) if oris else "", + "promoter_count": len(promoters), + "promoter_ids": ",".join(extract_ids(promoters)) if promoters else "", + "terminator_count": len(terminators), + "terminator_ids": ",".join(extract_ids(terminators)) if terminators else "", + "marker_count": len(markers), + "marker_ids": ",".join(extract_ids(markers)) if markers else "", + "cds_count": len(cdss), + "cds_ids": ",".join(extract_ids(cdss)) if cdss else "", + } + + +class Evaluator(EvalRunner): + """ + Evaluator that generates rollouts from a checkpoint and analyzes them. + + Loads prompts from CSV, generates samples using vLLM, analyzes each sequence + with plasmidkit, and returns a DataFrame with detailed annotation information. + """ + + def __init__(self, config: EvalConfig): + """ + Initialize the evaluator. + + Args: + config: Evaluation configuration + """ + self.config = config + self.llm: Optional[LLM] = None + self.base_config = Config() # For default prompts + self.analyzer = SequenceAnalyzer(config) + + def run_with_trainer(self, trainer: Any, wandb_run: Optional[Any] = None) -> pd.DataFrame: + """ + Run evaluation using the trainer's model directly (already loaded on GPU). + + Args: + trainer: Trainer instance with vLLM model already loaded + wandb_run: Optional wandb run object for logging + + Returns: + DataFrame with detailed annotation information for each sequence + """ + # Load prompts + prompts = self._load_prompts() + + if not prompts: + print("[Evaluator] Warning: No prompts loaded, returning empty DataFrame") + return pd.DataFrame() + + # Use trainer's vLLM instance directly + llm = trainer.llm if hasattr(trainer, 'llm') else None + if llm is None: + print("[Evaluator] Warning: Trainer does not have 'llm' attribute, cannot use in-memory model") + return pd.DataFrame() + + print(f"[Evaluator] Using trainer's vLLM instance directly (no model reload needed)") + self.llm = llm + + # Get sampling parameters + sampling_params = self.config.sampling_params + if sampling_params is None: + sampling_params = SamplingParams( + max_tokens=512, + temperature=0.8, + top_p=0.95, + top_k=0, + ) + + # Expand prompts for multiple samples per prompt + expanded_prompts = [] + prompt_indices = [] + for i, prompt in enumerate(prompts): + for _ in range(self.config.num_samples_per_prompt): + expanded_prompts.append(prompt) + prompt_indices.append(i) + + print(f"[Evaluator] Generating {len(expanded_prompts)} samples from {len(prompts)} prompts") + + # Generate rollouts + outputs = self.llm.generate(expanded_prompts, sampling_params) + + # Process results and analyze each sequence + records = [] + for idx, output in enumerate(outputs): + prompt = output.prompt + completion = output.outputs[0].text.replace(" ", "") + full = prompt + completion + + # Clean sequence for analysis (remove non-DNA characters) + cleaned_full = re.sub(r'[^ATCG]', '', full.upper()) + + # Analyze sequence + try: + analysis = self.analyzer.analyze(cleaned_full) + except Exception as e: + print(f"[Evaluator] Warning: Failed to analyze sequence {idx}: {e}") + analysis = { + "ori_count": 0, + "ori_ids": "", + "promoter_count": 0, + "promoter_ids": "", + "terminator_count": 0, + "terminator_ids": "", + "marker_count": 0, + "marker_ids": "", + "cds_count": 0, + "cds_ids": "", + } + + records.append({ + "prompt": prompt, + "response": completion, + "full": full, + "length": len(cleaned_full), + "completion_length": len(completion), + "full_length": len(full), + **analysis, + }) + + # Convert to DataFrame + df = pd.DataFrame(records) + + print(f"[Evaluator] Generated and analyzed {len(df)} rollouts") + + return df + + def _load_prompts(self) -> List[str]: + """ + Load prompts from CSV/parquet file. + + Returns: + List of prompt strings + """ + if not self.config.prompts_path: + print("[Evaluator] Warning: No prompts_path specified, using default prompts") + return [self._get_default_prompts()] + + prompts_path = self.config.prompts_path + + # Check if file exists + if not os.path.exists(prompts_path): + print(f"[Evaluator] Warning: Prompts file not found: {prompts_path}") + return [self._get_default_prompts()] + + try: + # Load based on file extension + if prompts_path.endswith('.parquet'): + df = pd.read_parquet(prompts_path) + elif prompts_path.endswith('.csv'): + df = pd.read_csv(prompts_path) + else: + print(f"[Evaluator] Warning: Unsupported file format: {prompts_path}") + return [self._get_default_prompts()] + + # Extract prompts column + if self.config.prompts_column not in df.columns: + print(f"[Evaluator] Warning: Column '{self.config.prompts_column}' not found in prompts file") + print(f"[Evaluator] Available columns: {list(df.columns)}") + return [self._get_default_prompts()] + + prompts = df[self.config.prompts_column].dropna().tolist() + prompts = [str(p).strip() for p in prompts if str(p).strip()] + + print(f"[Evaluator] Loaded {len(prompts)} prompts from {prompts_path}") + return prompts + + except Exception as e: + print(f"[Evaluator] Error loading prompts: {e}") + return [self._get_default_prompts()] + + def _get_default_prompts(self) -> str: + """Get default prompt if prompts file is not available.""" + # Use the default query from config (GFP cassette) + return self.base_config.default_query + + def _initialize_model(self, checkpoint_path: str) -> None: + """ + Initialize vLLM model from checkpoint path. + + Args: + checkpoint_path: Path to model checkpoint + """ + # Check if model is already initialized for this checkpoint + if self.llm is not None: + # For now, always reinitialize - could optimize later + pass + + print(f"[Evaluator] Loading model from checkpoint: {checkpoint_path}") + + try: + # Initialize vLLM with checkpoint path + # vLLM can load from local checkpoint directories + model_kwargs = { + "trust_remote_code": True, + } + + self.llm = LLM(model=checkpoint_path, **model_kwargs) + print(f"[Evaluator] Model loaded successfully") + + except Exception as e: + print(f"[Evaluator] Error loading model: {e}") + import traceback + traceback.print_exc() + raise \ No newline at end of file diff --git a/src/eval/eval_config.py b/src/eval/eval_config.py new file mode 100644 index 0000000..27d361b --- /dev/null +++ b/src/eval/eval_config.py @@ -0,0 +1,33 @@ +from pydantic import BaseModel, ConfigDict +from vllm import SamplingParams +from typing import Optional + +class EvalConfig(BaseModel): + """ + Configuration for evaluation analysis. + + Similar to RewardConfig but focused on detailed annotation extraction + rather than scoring. + """ + model_config = ConfigDict(arbitrary_types_allowed=True) + + # Model configuration + model_name: str + model_path: str + + # Prompts configuration + prompts_path: Optional[str] = None # Path to CSV/parquet file with prompts + prompts_column: str = "prompt" # Column name containing prompts + num_samples_per_prompt: int = 10 # Number of samples to generate per prompt + + # Annotation configuration (similar to RewardConfig) + overlap_merge_threshold: float = 0.8 # Overlap merge threshold for annotations + + # Generation configuration + sampling_params: Optional[SamplingParams] = None + + # Logging configuration + write_to_wandb: bool = False + wandb_project: Optional[str] = None + wandb_run_name: Optional[str] = None + diff --git a/src/rewards/bioinformatics/reward_config.py b/src/rewards/bioinformatics/reward_config.py index 5de9774..86ec753 100644 --- a/src/rewards/bioinformatics/reward_config.py +++ b/src/rewards/bioinformatics/reward_config.py @@ -12,7 +12,7 @@ class RewardConfig(BaseModel): length_reward_bonus: float = 0.5 # bonus multiplier for sequences in ideal length range location_aware: bool = True # reward sequences that are located in the correct location (e.g. promoter then cds then terminator) # Penalty factor applied when min/max constraints are violated (outside of range) - violation_penalty_factor: float = 0.5 + violation_penalty_factor: float = 1.0 # Deprecated - use length_reward_mode instead length_penalty: bool = False @@ -20,7 +20,7 @@ class RewardConfig(BaseModel): ori_min: int = 1 ori_max: int = 1 allowed_oris: Optional[List[str]] = None - ori_weight: float = 1.0 + ori_weight: float = 1.5 promoter_min: int = 1 promoter_max: int = 1 diff --git a/src/runners/grpo.py b/src/runners/grpo.py index 9184df0..fea171f 100644 --- a/src/runners/grpo.py +++ b/src/runners/grpo.py @@ -6,16 +6,21 @@ from src.rewards.bioinformatics.scorer import Scorer from src.rewards.bioinformatics.reward_config import RewardConfig from src.rewards.bioinformatics.logger import RewardComponentLogger +from src.eval.eval import Evaluator +from src.eval.eval_config import EvalConfig +from src.utils.training_utils import EvalCallback, test_checkpoint_directory_write +from vllm import SamplingParams import datetime from typing import List import wandb from concurrent.futures import ThreadPoolExecutor from threading import Lock import re +import os # Configuration cfg = Config() -run_name = f"grpo-{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" +run_name = f"grpo-production-{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" # Dataset loading def load_train_val_datasets(): @@ -52,17 +57,22 @@ def select_prompt_column(ds): "pad_token_id": tok.pad_token_id, } -# Training configuration +# Training configuration - use /s3 mount point with prefix path +checkpoint_dir = f"/s3/{cfg.checkpoints_path.rstrip('/')}/grpo-production/{run_name}" + +# Test checkpoint directory write access before proceeding +test_checkpoint_directory_write(checkpoint_dir) + args = GRPOConfig( model_init_kwargs=model_init_kwargs, - output_dir=f"/s3/checkpoints/verl-grpo/{run_name}", + output_dir=checkpoint_dir, # Training parameters num_train_epochs=20, - learning_rate=3e-6, + learning_rate=cfg.grpo_learning_rate, lr_scheduler_type="constant", warmup_ratio=0.0, - per_device_train_batch_size=16, + per_device_train_batch_size=cfg.grpo_per_device_train_batch_size, gradient_accumulation_steps=1, max_steps=-1, max_grad_norm=0.5, @@ -71,6 +81,7 @@ def select_prompt_column(ds): # Logging and checkpointing save_strategy="steps", save_steps=100, + save_total_limit=5, # Keep last 5 checkpoints logging_strategy="steps", logging_steps=1, report_to=["wandb"], @@ -78,15 +89,15 @@ def select_prompt_column(ds): # Evaluation do_eval=True, eval_strategy="steps", - eval_steps=100, + eval_steps=50, # Evaluate every 50 steps # Optimization bf16=torch.cuda.is_available(), gradient_checkpointing=False, # GRPO-specific - beta=1e-3, - epsilon=0.2, + beta=cfg.grpo_beta, + epsilon=cfg.grpo_epsilon, loss_type="bnpo", scale_rewards=True, mask_truncated_completions=False, @@ -95,10 +106,10 @@ def select_prompt_column(ds): # Generation parameters remove_unused_columns=False, max_prompt_length=1024, - num_generations=8, + num_generations=cfg.grpo_num_generations, max_completion_length=256, - temperature=0.95, - top_p=0.90, + temperature=cfg.grpo_temperature, + top_p=cfg.grpo_top_p, # vLLM configuration use_vllm=True, @@ -106,20 +117,31 @@ def select_prompt_column(ds): vllm_mode="colocate", ) -# Reward configuration +# Reward configuration - production parameters from sweep reward_config = RewardConfig( - punish_mode=False, - length_penalty=True, - min_length=1000, + punish_mode=True, # Use punish mode for better constraint learning + length_reward_mode=True, + min_length=2000, max_length=30000, + ideal_min_length=7000, + ideal_max_length=20000, + length_reward_bonus=0.7085046275614012, ori_min=1, ori_max=1, + ori_weight=1.0, promoter_min=1, promoter_max=5, + promoter_weight=1.0, terminator_min=0, terminator_max=2, + terminator_weight=0.5, marker_min=1, marker_max=2, + marker_weight=1.0, + cds_min=1, + cds_max=5, + cds_weight=1.0, + location_aware=True, ) # Initialize scorer and logger @@ -127,6 +149,27 @@ def select_prompt_column(ds): reward_logger = RewardComponentLogger(log_frequency=1) component_lock = Lock() +# Initialize evaluation callback +eval_config = EvalConfig( + model_name=cfg.model, + model_path=cfg.model, # Will be overridden by checkpoint path in callback + prompts_path=cfg.val_dataset, # Use test.parquet for evaluation prompts + prompts_column="prompt", + num_samples_per_prompt=5, # Fewer samples for quick testing + overlap_merge_threshold=0.8, + sampling_params=SamplingParams( + max_tokens=256, + temperature=0.95, + top_p=0.90, + top_k=0, + ), + write_to_wandb=True, + wandb_project=cfg.wandb_project, + wandb_run_name=run_name, +) +evaluator = Evaluator(eval_config) +eval_callback = EvalCallback(evaluator) + # Reward function def score_single(idx_and_seq): """Score a single sequence and log components thread-safely.""" @@ -161,28 +204,38 @@ def batch_reward_fn(prompts: List[str], completions: List[str], **kwargs) -> Lis return [r[0] for r in results] # Initialize W&B -wandb.init( +wandb_run = wandb.init( project=cfg.wandb_project, entity=cfg.wandb_entity, name=run_name, + tags=["production", "grpo", "optimized-hyperparams"], config={ "model": cfg.model, "reward_config": reward_config.model_dump(), "training": { - "learning_rate": args.learning_rate, - "batch_size": args.per_device_train_batch_size, + "learning_rate": cfg.grpo_learning_rate, + "batch_size": cfg.grpo_per_device_train_batch_size, "num_epochs": args.num_train_epochs, + "num_generations": cfg.grpo_num_generations, }, "grpo": { - "beta": args.beta, - "epsilon": args.epsilon, - "temperature": args.temperature, - "num_generations": args.num_generations, + "beta": cfg.grpo_beta, + "epsilon": cfg.grpo_epsilon, + "temperature": cfg.grpo_temperature, + "top_p": cfg.grpo_top_p, "loss_type": args.loss_type, }, + "checkpoint_dir": checkpoint_dir, }, ) +# Print wandb URL and checkpoint info +if wandb_run: + print(f"\n{'='*80}") + print(f"🚀 W&B Run URL: {wandb_run.url}") + print(f"📁 Checkpoint Directory: {checkpoint_dir}") + print(f"{'='*80}\n") + # Initialize trainer trainer = GRPOTrainer( model=cfg.model, @@ -191,10 +244,32 @@ def batch_reward_fn(prompts: List[str], completions: List[str], **kwargs) -> Lis train_dataset=train_ds, eval_dataset=eval_ds, processing_class=tok, - callbacks=[reward_logger], + callbacks=[reward_logger, eval_callback], ) +# Set trainer reference in callback (for accessing trainer.llm) +eval_callback.set_trainer(trainer) + # Train and save +print(f"Starting training with {args.num_train_epochs} epochs...") trainer.train() -trainer.save_model(args.output_dir) -tok.save_pretrained(args.output_dir) + +# Save final model and tokenizer +print(f"Saving final model to {checkpoint_dir}...") +trainer.save_model(checkpoint_dir) +tok.save_pretrained(checkpoint_dir) + +# Log final checkpoint as W&B artifact +artifact = wandb.Artifact( + name=f"model-{run_name}", + type="model", + description=f"Final GRPO model checkpoint from production run with optimized hyperparameters", +) +artifact.add_dir(checkpoint_dir) +wandb_run.log_artifact(artifact) + +print(f"✓ Training complete! Model saved to {checkpoint_dir}") +print(f"✓ Model artifact logged to W&B: {wandb_run.url}") + +# Finish run +wandb.finish() diff --git a/src/runners/grpo_sweep.py b/src/runners/grpo_sweep.py index 44d2ec3..6f95db3 100644 --- a/src/runners/grpo_sweep.py +++ b/src/runners/grpo_sweep.py @@ -11,12 +11,17 @@ from src.rewards.bioinformatics.scorer import Scorer from src.rewards.bioinformatics.reward_config import RewardConfig from src.rewards.bioinformatics.logger import RewardComponentLogger +from src.eval.eval import Evaluator +from src.eval.eval_config import EvalConfig +from src.utils.training_utils import EvalCallback, test_checkpoint_directory_write +from vllm import SamplingParams import datetime from typing import List import wandb from concurrent.futures import ThreadPoolExecutor from threading import Lock import re +import os def load_train_val_datasets(cfg): @@ -39,7 +44,11 @@ def main(): # Initialize W&B run with a meaningful name run_name = f"grpo-sweep-{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" - run = wandb.init(name=run_name) + run = wandb.init( + name=run_name, + entity=cfg.wandb_entity, + project=cfg.wandb_project, + ) # Get sweep params sweep_config = wandb.config @@ -65,8 +74,12 @@ def main(): "pad_token_id": tok.pad_token_id, } - # Training configuration with sweep parameters - checkpoint_dir = f"/mnt/s3/phd-research-storage-1758274488/checkpoints/grpo-sweeps/{run_name}" + # Training configuration - use /s3 mount point with prefix path + checkpoint_dir = f"/s3/{cfg.checkpoints_path.rstrip('/')}/grpo-sweeps/{run_name}" + + # Test checkpoint directory write access before proceeding + test_checkpoint_directory_write(checkpoint_dir) + args = GRPOConfig( model_init_kwargs=model_init_kwargs, output_dir=checkpoint_dir, @@ -93,7 +106,7 @@ def main(): # Evaluation do_eval=True, eval_strategy="steps", - eval_steps=100, + eval_steps=50, # More frequent eval for sweeps to track progress # Optimization bf16=torch.cuda.is_available(), @@ -148,11 +161,35 @@ def main(): location_aware=sweep_config.get("reward_location_aware", True), ) + # Log reward_config to wandb + wandb.config.update({"reward_config": reward_config.model_dump()}) + # Initialize scorer and logger scorer = Scorer(reward_config) reward_logger = RewardComponentLogger(log_frequency=5) # Log frequently for short runs component_lock = Lock() + # Initialize evaluation callback + eval_config = EvalConfig( + model_name=cfg.model, + model_path=cfg.model, + prompts_path=cfg.val_dataset, # Use test.parquet for evaluation prompts + prompts_column="prompt", + num_samples_per_prompt=5, # Fewer samples for sweeps + overlap_merge_threshold=0.8, + sampling_params=SamplingParams( + max_tokens=256, + temperature=0.95, + top_p=0.90, + top_k=0, + ), + write_to_wandb=True, + wandb_project=cfg.wandb_project, + wandb_run_name=run_name, + ) + evaluator = Evaluator(eval_config) + eval_callback = EvalCallback(evaluator) + # Reward function def score_single(idx_and_seq): idx, seq = idx_and_seq @@ -179,9 +216,12 @@ def batch_reward_fn(prompts: List[str], completions: List[str], **kwargs) -> Lis train_dataset=train_ds, eval_dataset=eval_ds, processing_class=tok, - callbacks=[reward_logger], + callbacks=[reward_logger, eval_callback], ) + # Set trainer reference in callback (for accessing trainer.llm) + eval_callback.set_trainer(trainer) + # Train trainer.train() diff --git a/src/utils/training_utils.py b/src/utils/training_utils.py new file mode 100644 index 0000000..91af73f --- /dev/null +++ b/src/utils/training_utils.py @@ -0,0 +1,243 @@ +from transformers import TrainerCallback +from typing import Any, Protocol, Optional +import wandb +import pandas as pd +import os +import sys +import datetime +from abc import ABC, abstractmethod + + +def test_checkpoint_directory_write(checkpoint_dir: str) -> None: + """ + Test that checkpoint directory exists and is writable. + Raises an error and exits if write test fails. + + Args: + checkpoint_dir: Path to checkpoint directory to test + """ + try: + # Create directory if it doesn't exist + os.makedirs(checkpoint_dir, exist_ok=True) + + # Test write access + test_file = os.path.join(checkpoint_dir, ".write_test") + test_content = f"Write test at {datetime.datetime.now().isoformat()}\n" + + # Write test file + with open(test_file, 'w') as f: + f.write(test_content) + + # Read back to verify + with open(test_file, 'r') as f: + read_content = f.read() + + if read_content != test_content: + raise IOError(f"Write test failed: content mismatch") + + # Clean up test file + os.remove(test_file) + + print(f"✓ Checkpoint directory write test passed: {checkpoint_dir}") + + except Exception as e: + error_msg = ( + f"\n{'='*80}\n" + f"❌ CHECKPOINT DIRECTORY WRITE TEST FAILED\n" + f"{'='*80}\n" + f"Directory: {checkpoint_dir}\n" + f"Error: {str(e)}\n" + f"\nThis usually means:\n" + f" - S3 mount is not available at /s3\n" + f" - Insufficient permissions to write to S3\n" + f" - Disk space is full\n" + f"\nPlease check:\n" + f" 1. Docker volume mount: /mnt/s3/phd-research-storage-1758274488:/s3:rw\n" + f" 2. S3 mount is accessible: ls -la /s3\n" + f" 3. Write permissions on S3 mount\n" + f"{'='*80}\n" + ) + print(error_msg, file=sys.stderr) + sys.exit(1) + + +class EvalRunner(ABC): + """ + Protocol for evaluation runner classes. + + Classes implementing this protocol should have a `run_with_trainer()` method that: + - Takes trainer and wandb_run as parameters + - Uses the trainer's model directly (already loaded on GPU) + - Returns a pandas DataFrame containing evaluation results + - Each row should contain evaluation metrics (e.g., score, length, etc.) + """ + + @abstractmethod + def run_with_trainer(self, trainer: Any, wandb_run: Optional[Any] = None) -> pd.DataFrame: + """ + Run evaluation using the trainer's model directly. + + Args: + trainer: Trainer instance with model already loaded + wandb_run: Optional wandb run object for logging + + Returns: + DataFrame with evaluation results, each row containing evaluation metrics + """ + ... + + +class EvalCallback(TrainerCallback): + """ + Minimal training callback that runs evaluation during training and logs results to wandb. + + Uses the trainer's in-memory model directly (no need to reload from disk). + Delegates all evaluation logic to the provided evaluator class. + """ + + def __init__(self, evaluator: EvalRunner): + """ + Initialize the evaluation callback. + + Args: + evaluator: An object with a `run_with_trainer(trainer, wandb_run)` method that + performs evaluation using the trainer's model and returns a DataFrame + """ + self.evaluator = evaluator + self.last_eval_step = -1 + self._trainer_ref = None # Will be set when trainer is available + + def set_trainer(self, trainer: Any): + """Set trainer reference for callbacks that need it.""" + self._trainer_ref = trainer + + def on_evaluate(self, args, state, control, **kwargs): + """Run evaluation when trainer evaluation is triggered.""" + print(f"[EvalCallback] on_evaluate called at step {state.global_step}") + + # Avoid duplicate evals at the same step + if state.global_step == self.last_eval_step: + print(f"[EvalCallback] Already evaluated at step {state.global_step}, skipping") + return + + self.last_eval_step = state.global_step + + try: + # Get trainer from kwargs - transformers TrainerCallback passes 'model' and sometimes 'trainer' + # For GRPOTrainer, we need to access it via the callback's parent reference or kwargs + trainer = kwargs.get('trainer') + if trainer is None: + # Try to get model directly - GRPOTrainer might pass model in kwargs + model = kwargs.get('model') + if model is not None: + # If we have model but not trainer, we need to find trainer another way + # Actually, let's check if we can store a reference to trainer in __init__ + print("[EvalCallback] Warning: Trainer not found in kwargs, checking callback context") + # Fallback: use self if callback was passed trainer reference + if hasattr(self, '_trainer_ref'): + trainer = self._trainer_ref + else: + print("[EvalCallback] Warning: Cannot access trainer, skipping evaluation") + return + + print(f"[EvalCallback] Running evaluation at step {state.global_step} (using model from trainer)") + + # Get wandb run object and URL + wandb_run = wandb.run + if wandb_run: + wandb_url = wandb_run.url + print(f"[EvalCallback] W&B Run URL: {wandb_url}") + + # Run evaluation using the trainer's model directly + results_df = self.evaluator.run_with_trainer(trainer, wandb_run) + + if results_df is None or len(results_df) == 0: + print("[EvalCallback] Warning: Evaluation returned no results") + return + + # Log results to wandb + self._log_results(results_df, state.global_step) + + print(f"[EvalCallback] Logged evaluation results for step {state.global_step}") + + except Exception as e: + print(f"[EvalCallback] Error during evaluation: {e}") + import traceback + traceback.print_exc() + + def _get_checkpoint_path(self, args, state) -> Optional[str]: + """ + Get the path to the current checkpoint. + + Args: + args: Training arguments + state: Trainer state + + Returns: + Path to checkpoint directory, or base model path if checkpoint doesn't exist yet + """ + # Check if there's a checkpoint directory for this step + checkpoint_dir = f"{args.output_dir}/checkpoint-{state.global_step}" + if os.path.exists(checkpoint_dir) and os.path.exists(f"{checkpoint_dir}/config.json"): + return checkpoint_dir + + # Fallback: use output_dir if checkpoint doesn't exist yet + # Check if output_dir has config.json (meaning it's a valid checkpoint) + if os.path.exists(args.output_dir) and os.path.exists(f"{args.output_dir}/config.json"): + return args.output_dir + + # If no checkpoint exists yet, return None - evaluator will use base model + return None + + def _log_results(self, results_df: pd.DataFrame, step: int) -> None: + """ + Log evaluation results to wandb as both table and artifact. + + Args: + results_df: DataFrame containing evaluation results + step: Current training step + """ + if results_df is None or len(results_df) == 0: + return + + df = results_df + + # Log as wandb table (for quick dashboard viewing) + wandb.log({ + "eval/step": step, + "eval/results_table": wandb.Table(dataframe=df), + }) + + # Log summary statistics for numeric columns + numeric_cols = df.select_dtypes(include=['number']).columns + stats = {} + for col in numeric_cols: + if col != "step": + stats[f"eval/stats/{col}/mean"] = float(df[col].mean()) + stats[f"eval/stats/{col}/std"] = float(df[col].std()) + stats[f"eval/stats/{col}/min"] = float(df[col].min()) + stats[f"eval/stats/{col}/max"] = float(df[col].max()) + stats[f"eval/stats/{col}/median"] = float(df[col].median()) + + if stats: + wandb.log(stats) + + # Create and log artifact (for full data export) + artifact_name = f"eval_results_step_{step}" + artifact = wandb.Artifact( + name=artifact_name, + type="evaluation_results", + description=f"Evaluation results at training step {step}", + metadata={ + "step": step, + "total_samples": len(df), + } + ) + + # Add CSV to artifact + with artifact.new_file("results.csv", mode="w") as f: + df.to_csv(f, index=False) + + # Log artifact + wandb.log_artifact(artifact) + diff --git a/sweeps/configs/sweep_config_training_with_eval.yaml b/sweeps/configs/sweep_config_training_with_eval.yaml new file mode 100644 index 0000000..61f1bfe --- /dev/null +++ b/sweeps/configs/sweep_config_training_with_eval.yaml @@ -0,0 +1,144 @@ +# W&B Sweep Configuration for PlasmidRL GRPO - Training + Length Reward + Eval +# Focus: Broad training hyperparameter search with length-based rewards and evaluation +# Duration: 500 steps per trial for stable evaluation +# Includes: Evaluation callback to track detailed sequence analysis +# +# Strategy: Test 2 length configurations across full hyperparameter space +# +# Usage: +# 1. Initialize the sweep: +# wandb sweep sweeps/configs/sweep_config_training_with_eval.yaml +# +# 2. Copy the sweep ID from output (e.g., mcclain/plasmidrl-grpo-sweeps/abc123xyz) +# +# 3. Run agent(s): +# SWEEP_ID=mcclain/plasmidrl-grpo-sweeps/abc123xyz docker compose up grpo-sweep +# +# Monitor progress at: https://wandb.ai/mcclain/plasmidrl-grpo-sweeps + +program: sweeps/run_sweep_agent.py +method: bayes +metric: + name: reward_components/total_reward/mean + goal: maximize + +parameters: + # ==================== TRAINING HYPERPARAMETERS (BROAD SEARCH) ==================== + + max_steps: + value: 500 # Longer runs for stable evaluation + + # Learning rate - broad exploration + learning_rate: + distribution: log_uniform_values + min: 1e-6 + max: 1e-4 + + # Batch size + per_device_train_batch_size: + values: [8, 16, 32] + + # Generations per batch + num_generations: + values: [4, 8, 16] + + # Temperature - sampling randomness + temperature: + distribution: uniform + min: 0.7 + max: 1.4 + + # Top-p - nucleus sampling + top_p: + distribution: uniform + min: 0.85 + max: 0.95 + + # Beta - KL penalty coefficient + beta: + distribution: log_uniform_values + min: 1e-4 + max: 1e-2 + + # Epsilon - PPO-style clipping + epsilon: + distribution: uniform + min: 0.1 + max: 0.3 + + # ==================== LENGTH REWARD PARAMETERS (2 COMBOS) ==================== + + reward_length_reward_mode: + value: true # Enable length-based rewards + + # Test 2 length configurations + # Combo 1: min=2000, ideal_min=3000, ideal_max=12000, max=15000 + # Combo 2: min=5000, ideal_min=7000, ideal_max=20000, max=30000 + + reward_min_length: + values: [2000, 5000] + + reward_max_length: + values: [15000, 30000] + + reward_ideal_min_length: + values: [3000, 7000] + + reward_ideal_max_length: + values: [12000, 20000] + + # Bonus multiplier for being in ideal range + reward_length_reward_bonus: + distribution: uniform + min: 0.3 + max: 0.8 + + # ==================== REWARD CONFIG (FIXED - STANDARD) ==================== + + reward_punish_mode: + value: true + + reward_promoter_max: + value: 5 + + reward_terminator_max: + value: 2 + + reward_marker_max: + value: 2 + + reward_cds_max: + value: 5 + + reward_location_aware: + value: true + + # Reward component weights (all enabled with standard weights) + reward_ori_weight: + value: 1.0 + + reward_promoter_weight: + value: 1.0 + + reward_terminator_weight: + value: 0.5 + + reward_marker_weight: + value: 1.0 + + reward_cds_weight: + value: 1.0 + +# Early termination - adjusted for longer runs +early_terminate: + type: hyperband + min_iter: 100 # Higher threshold for 500-step runs + eta: 3 + s: 2 + +# Optional: Limit total number of runs +# run_cap: 50 + + + +