inference.py

import argparse
import json
import pathlib
import random
from typing import Any, Dict, List, Optional, Tuple

import torch
import torch.nn.functional as F
from tqdm import tqdm
import sys
sys.path.append(str(pathlib.Path(__file__).resolve().parent.parent))  # Add project root to path for imports
sys.path.append('src.models')

from src.utils.inference_conditions import build_conditions
from src.utils.inference_generation_utils import (
    build_cas13_full_sequence,
    build_initial_tokens,
    calculate_seq_similarity,
    collect_diffusion_func_cond,
    extract_guide_seq_from_full,
    normalize_score_output,
    parse_design_pos,
    prepare_ss_extra_state,
)
from src.utils.inference_models import load_diffusion_model, load_guidance_model
from src.utils.inference_targets import (
    build_task_targets_from_args,
    load_data_records,
    merge_args_with_record,
    parse_json_map,
    parse_task_list,
)
from src.utils.inference_types import GuidanceConfig, MODEL_SEQUENCE_LENGTH
from src.utils.rna_fm_tokenizer import RNAFMAlphabetTokenizer

import os

def run_multi_task_guided_inference(
    *,
    tasks: List[str],
    task_targets: Dict[str, Any],
    task_scales: Dict[str, float],
    tokenizer,
    lens: int,
    diffusion_model,
    guidance_models: Dict[str, Any],
    conditions,
    guidance: GuidanceConfig,
    template: Optional[str],
    seed: Optional[str],
    seed_mask_num: int,
    mask_ratio: float,
    design_pos: Optional[List[Tuple[int, int]]] = None,
    max_lens: int = 1024,
    seq_cond_enabled: bool = True,
) -> Dict[str, Any]:
    device = next(diffusion_model.parameters()).device
    uncon_only = len(tasks) == 1 and tasks[0] == "uncon"

    rna_ids, seq_len = build_initial_tokens(
        tokenizer,
        lens=lens,
        template=template,
        seed=seed,
        seed_mask_num=seed_mask_num,
        mask_ratio=mask_ratio,
        device=device,
        design_pos=design_pos,
    )
    
    rna_ids = rna_ids[:, :seq_len]

    seq_cond = None
    if seq_cond_enabled:
        if conditions.ts.seq_cond is not None :
            seq_cond = conditions.ts.seq_cond[:, :seq_len]
        elif conditions.ss.seq_cond is not None:
            seq_cond = conditions.ss.seq_cond[:, :seq_len]
        elif conditions.proteins.seq_cond is not None:
            seq_cond = conditions.proteins.seq_cond[:, :seq_len]
    else:
        conditions.ts.seq_cond, conditions.ss.seq_cond, conditions.proteins.seq_cond = None, None, None

    if seq_cond is not None:
        seq_cond = seq_cond.to(device)

    func_cond = None if uncon_only else collect_diffusion_func_cond(conditions)

    label_cond = None
    if (not uncon_only) and "rfam" in tasks and conditions.rfam.label is not None:
        label_cond = torch.tensor([conditions.rfam.label], dtype=torch.long, device=device)

    constrain_mask = (rna_ids != tokenizer.mask_id) if guidance.constrain_logits else None
    rna_input_mask = rna_ids.ne(tokenizer.pad_id).long()
    rna_target_mask = torch.zeros_like(rna_ids).bool()

    ss_extra_state = {} if uncon_only else prepare_ss_extra_state(tasks, conditions, tokenizer, device)

    best_combined = -1e9
    best_rna_ids = rna_ids.clone()

    masked_positions = [i for i in range(seq_len) if rna_ids[0, i] == tokenizer.mask_id]
    order = masked_positions.copy()
    random.shuffle(order)

    D = len(masked_positions)

    for step, pos in enumerate(order):
        t_val = D - (len(order) - step) + 1
        timesteps = torch.tensor([t_val], dtype=torch.long, device=device)

        logits = diffusion_model.model(
            rna_ids=rna_ids,
            twod_tokens=None,
            timestep=timesteps,
            func_cond=func_cond,
            label_cond=label_cond,
            rna_input_mask=rna_input_mask,
            seq_cond=seq_cond,
        )

        if constrain_mask is not None:
            batch_idx, pos_idx = torch.where(constrain_mask)
            token_idx = rna_ids[batch_idx, pos_idx]
            logits[batch_idx, pos_idx, token_idx] = 1e9

        guided_logits = logits
        step_scores: Dict[str, float] = {}

        if guidance.update_by_guide:
            for task in tasks:
                guidance_model = guidance_models.get(task)
                if guidance_model is None:
                    continue
                guidance_model.train()
                target = task_targets.get(task)
                scale = float(task_scales.get(task, guidance.guidance_scale))
                guided_logits, score = guidance_model.apply_guidance(
                    logits=guided_logits,
                    rna_ids=rna_ids,
                    pos=pos,
                    y_target=target,
                    guidance_scale=scale,
                    rna_input_mask=rna_input_mask,
                    conditions=conditions,
                    update_all=guidance.update_all,
                    use_gumbel=guidance.use_gumbel,
                    tokenizer=tokenizer,
                    extra_state=ss_extra_state if task == "ss" else None,
                )
                step_scores[task] = float(score)

        probs = F.softmax(guided_logits, dim=-1)
        if not guidance.update_all:
            new_token = probs[0, pos].argmax(-1) if guidance.argmax else torch.multinomial(probs[0, pos], 1).item()
            rna_ids[0, pos] = new_token
        else:
            rna_ids[rna_target_mask] = probs.argmax(-1)[rna_target_mask]
            rna_ids[0, pos] = probs[0, pos].argmax(-1) if guidance.argmax else torch.multinomial(probs[0, pos], 1)
        rna_target_mask[0, pos] = True

        if step_scores:
            combined = sum(step_scores.values()) / len(step_scores)
            if combined > best_combined and guidance.use_best:
                best_combined = combined
                best_rna_ids = logits.argmax(-1).clone()

    final_rna_ids = rna_ids
    final_tokens = final_rna_ids[0].detach().cpu().numpy()
    final_seq = tokenizer.detokenize(final_tokens)
    final_scores: Dict[str, float] = {}
    best_scores: Dict[str, float] = {}
    final_accs: Dict[str, float] = {}
    best_accs: Dict[str, float] = {}

    for task in tasks:
        guidance_model = guidance_models.get(task)
        if guidance_model is None:
            continue
        guidance_model.eval()
        target = task_targets.get(task)

        final_score_raw = guidance_model.score(
            rna_ids=final_rna_ids,
            y_target=target,
            rna_input_mask=rna_input_mask,
            conditions=conditions,
            tokenizer=tokenizer,
            extra_state=ss_extra_state if task == "ss" else None,
        )
        final_metric = normalize_score_output(final_score_raw)
        final_scores[task] = final_metric["score"]
        final_accs[task] = final_metric["acc"]

        best_score_raw = guidance_model.score(
            rna_ids=best_rna_ids,
            y_target=target,
            rna_input_mask=rna_input_mask,
            conditions=conditions,
            tokenizer=tokenizer,
            extra_state=ss_extra_state if task == "ss" else None,
        )
        best_metric = normalize_score_output(best_score_raw)
        best_scores[task] = best_metric["score"]
        best_accs[task] = best_metric["acc"]

    final_acc_mean = float(sum(final_accs.values()) / len(final_accs)) if final_accs else 0.0
    best_acc_mean = float(sum(best_accs.values()) / len(best_accs)) if best_accs else 0.0
    final_score_mean = float(sum(final_scores.values()) / len(final_scores)) if final_scores else 0.0
    best_score_mean = float(sum(best_scores.values()) / len(best_scores)) if best_scores else 0.0

    use_best = (best_acc_mean > final_acc_mean) or (
        best_acc_mean == final_acc_mean and best_score_mean > final_score_mean
    )

    selected_tokens = best_rna_ids[0].detach().cpu().numpy() if use_best else final_tokens
    selected_seq = tokenizer.detokenize(selected_tokens)
    selected_task_scores = best_scores if use_best else final_scores
    selected_task_accs = best_accs if use_best else final_accs

    if "ts" in tasks:
        selected_task_scores = dict(selected_task_scores)
        selected_task_scores["ts"] = 0.0
        selected_task_accs = dict(selected_task_accs)
        ts_acc = 0.0
        ref_seq = conditions.ts.reference_sequence
        if ref_seq:
            ts_acc = calculate_seq_similarity(selected_seq, ref_seq)
        selected_task_accs["ts"] = ts_acc

    selected_acc_mean = float(sum(selected_task_accs.values()) / len(selected_task_accs)) if selected_task_accs else 0.0

    return {
        "sequence": selected_seq,
        "final_sequence": final_seq,
        "length": len(selected_seq),
        "task_scores": selected_task_scores,
        "task_acc": selected_task_accs,
        "selected_acc": selected_acc_mean,
        "selected_from": "best_seq" if use_best else "final_seq",
        "active_tasks": tasks,
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Unified multi-task inference with task-specific guidance_models")
    parser.add_argument("--tasks", type=str, default="uncon", help="Comma-separated tasks: uncon,text,rfam,disease,prot_rna_bind,ss,ts,utr,cas13")
    parser.add_argument("--config", type=str, default="./configs/rna_oadm_config.json")
    parser.add_argument("--diffusion_model_path", type=str, default="checkpoints/pretrained_gen.pt", help="Path to the main diffusion model checkpoint")

    parser.add_argument("--guidance_model_path", type=str, default="", help="Single-task guidance model path (uses default if omitted)")
    parser.add_argument("--guidance_model_paths", type=str, default="{}", help='JSON map, e.g. {"rfam":"path/to/rfam.pt"}')
    parser.add_argument("--task_targets", type=str, default="{}", help='JSON fallback for int targets, e.g. {"rfam":12}')
    parser.add_argument("--guidance_scales", type=str, default="{}", help='JSON map, e.g. {"rfam":1.5,"ss":2.0}')

    parser.add_argument("--rfam_target", type=str, default="RF00001", help="RFAM family ID, e.g. RF00001. Auto-maps to int label.")
    parser.add_argument("--disease_target", type=str, default="", help="Disease name, e.g. 'Osteoarthritis'. Auto-maps to int label.")
    parser.add_argument("--protein_targets", type=str, default="CAPRIN1,-AARS", help="Comma-separated protein names, e.g. 'CAPRIN1,AGO'. Auto-maps to int indices.")

    parser.add_argument("--ernie_pretrained", type=str, default="checkpoints/ERNIE-RNA_pretrain.pt")
    parser.add_argument("--ernie_ss_ckpt", type=str, default="checkpoints/ERNIE-RNA_attn-map_ss_prediction_bpRNA-1m-all_and_RNAStralign_checkpoint.pt")

    parser.add_argument("--lens", type=int, default=101)
    parser.add_argument("--num_samples", type=int, default=10)
    parser.add_argument("--template", type=str, default="")
    parser.add_argument("--seed", type=str, default="", help="Seed sequence for partial generation (only for miRNA generation, ignored if --sequence is provided)")
    parser.add_argument("--seed_mask_num", type=int, default=0)
    parser.add_argument("--mask_ratio", type=float, default=1.0)
    parser.add_argument("--design_pos", type=str, default="", help="Design positions in format 'start-end' for single range or '(start1-end1,start2-end2,...)'")

    parser.add_argument("--target_before", type=str, default="", help="Target context before guide region (20bp for CAS13)")
    parser.add_argument("--target_at_guide", type=str, default="", help="Target region at guide site (28bp for CAS13)")
    parser.add_argument("--target_after", type=str, default="", help="Target context after guide region (20bp for CAS13)")

    parser.add_argument("--func_cond", type=str, default="RNA aptamer with high stability", help="Plain-text condition used by task 'text'")
    parser.add_argument("--rfam_func_cond", type=str, default="")
    parser.add_argument("--protein_func_cond", type=str, default="")
    parser.add_argument("--disease_func_cond", type=str, default="")
    parser.add_argument("--secondary_structure", type=str, default="")
    parser.add_argument("--seq_cond_enable", type=bool, default=True)
    parser.add_argument("--icshape", type=str, default="")
    parser.add_argument("--pdb_path", type=str, default="", help="PDB file path for ts (3D structure-conditioned) task")
    parser.add_argument("--structure_encoder_root", type=str, default="src/structure_encoder", help="Path to structure encoder project root")
    parser.add_argument("--structure_encoder_ckpt", type=str, default="checkpoints/structure_encoder.pt", help="Path to structure encoder checkpoint")

    parser.add_argument("--guidance_scale", type=float, default=1.0)
    parser.add_argument("--update_all", action="store_true", default=True)
    parser.add_argument("--use_gumbel", action="store_true", default=False)
    parser.add_argument("--argmax", action="store_true", default=False)
    parser.add_argument("--constrain_logits", action="store_true", default=True)
    parser.add_argument("--update_by_guide", action="store_true", default=True)
    parser.add_argument("--use_best", action="store_true", default=True)

    parser.add_argument("--data_path", type=str, default="", help="Path to a JSON file whose top-level value is a list of per-item inputs")
    parser.add_argument("--output_path", type=str, default="results/infer_results/results.json", help="Path to save the inference results JSON  ")
    return parser


def main() -> None:
    args = build_parser().parse_args()

    with open(args.config, "r", encoding="utf-8") as f:
        config = json.load(f)

    tokenizer = RNAFMAlphabetTokenizer()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    tasks = parse_task_list(args.tasks)
    guidance_model_paths = parse_json_map(args.guidance_model_paths)
    task_targets_raw = parse_json_map(args.task_targets)
    task_scales = parse_json_map(args.guidance_scales)

    diffusion_model_path = args.diffusion_model_path
    default_diffusion_model_path = build_parser().get_default("diffusion_model_path")
    project_root = pathlib.Path(__file__).resolve().parent
    if "ts" in tasks:
        config["seq_cond_dim"] = 128
        config["seq_cond_enabled"] = True
        diffusion_model_path = str(project_root / "checkpoints" / "ts_gen.pt")
    elif "utr" in tasks:
        config["seq_cond_enabled"] = False
        config["func_cond_enabled"] = False
        diffusion_model_path = str(project_root / "checkpoints" / "utr_gen.pt")
    elif "cas13" in tasks:
        config["seq_cond_enabled"] = False
        config["func_cond_enabled"] = False
        diffusion_model_path = str(project_root / "checkpoints" / "cas13_gen.pt")
    elif "text" in tasks:
        config["seq_cond_enabled"] = False
        config["func_cond_enabled"] = True
        if args.diffusion_model_path == default_diffusion_model_path:
            diffusion_model_path = str(project_root / "checkpoints" / "pretrained_gen.pt")

    if len(tasks) == 1:
        task = tasks[0]
        if task not in guidance_model_paths and args.guidance_model_path:
            guidance_model_paths[task] = args.guidance_model_path

    guidance = GuidanceConfig(
        guidance_scale=args.guidance_scale,
        update_all=args.update_all,
        use_gumbel=args.use_gumbel,
        update_by_guide=args.update_by_guide,
        argmax=args.argmax,
        constrain_logits=args.constrain_logits,
        use_best=args.use_best,
    )

    diffusion_model = load_diffusion_model(
        config=config,
        tokenizer=tokenizer,
        ckpt_path=diffusion_model_path,
        ex_name="multi_task_diffusion",
    )

    guidance_models: Dict[str, Any] = {}
    for task in tasks:
        if (not guidance.update_by_guide) or task in {"uncon", "ts", "text"}:
            guidance_models[task] = None
            continue

        ckpt_path = guidance_model_paths.get(task)
        guidance_models[task] = load_guidance_model(
            task=task,
            config=config,
            tokenizer=tokenizer,
            ckpt_path=ckpt_path,
            args=args,
        )

    records = load_data_records(args.data_path)
    if not records:
        records = [None]

    results = []

    total_records = len(records)
    for record_idx, record in enumerate(tqdm(records, desc="Records"), start=1):
        record_args = merge_args_with_record(args, record)
        task_targets, _ = build_task_targets_from_args(record_args, task_targets_raw, tasks)
        conditions = build_conditions(record_args, device, task_targets, tasks)

        run_count = 1 if args.data_path else args.num_samples
        for sample_idx, _ in enumerate(tqdm(range(run_count), desc="Inference"), start=1):
            infer_lens = record_args.lens
            template = record_args.template.replace("T", "U") if record_args.template else None

            design_pos_str = getattr(record_args, "design_pos", "") or ""
            design_pos: Optional[List[Tuple[int, int]]] = None

            if "cas13" in tasks and conditions.cas13.target_before and conditions.cas13.target_at_guide and conditions.cas13.target_after:
                full_seq = build_cas13_full_sequence(
                    target_before=conditions.cas13.target_before,
                    target_at_guide=conditions.cas13.target_at_guide,
                    target_after=conditions.cas13.target_after,
                    guide_seq=conditions.cas13.guide_seq if conditions.cas13.guide_seq else "N" * 28,
                )
                template = full_seq.upper().replace("T", "U")
                infer_lens = MODEL_SEQUENCE_LENGTH
                if not design_pos_str:
                    design_pos_str = "88-115"

            if template is not None and design_pos_str:
                try:
                    design_pos = parse_design_pos(design_pos_str)
                except ValueError as e:
                    print(f"Warning: Invalid design_pos '{design_pos_str}': {e}. Ignoring design_pos.")
                    design_pos = None

            if "ts" in tasks and conditions.ts.seq_cond is not None and not record_args.template:
                infer_lens = int(conditions.ts.seq_cond.shape[1])
            if "ss" in tasks and conditions.ss.seq_cond is not None and not record_args.template:
                infer_lens = int(conditions.ss.seq_cond.shape[1])

            result = run_multi_task_guided_inference(
                tasks=tasks,
                task_targets=task_targets,
                task_scales={k: float(v) for k, v in task_scales.items()},
                tokenizer=tokenizer,
                lens=infer_lens,
                diffusion_model=diffusion_model,
                guidance_models=guidance_models,
                conditions=conditions,
                guidance=guidance,
                template=template,
                seed=record_args.seed if record_args.seed else None,
                seed_mask_num=record_args.seed_mask_num,
                mask_ratio=record_args.mask_ratio,
                design_pos=design_pos,
                seq_cond_enabled=record_args.seq_cond_enable,
            )

            if "cas13" in tasks and template and len(template) == MODEL_SEQUENCE_LENGTH:
                guide_seq_only = extract_guide_seq_from_full(result["sequence"])
                result["guide_seq"] = guide_seq_only
                result["full_sequence"] = result["sequence"]
                result["sequence"] = guide_seq_only
                result["design_pos"] = design_pos_str

            if "ts" in tasks and conditions.ts.pdb_path is not None:
                result["pdb_path"] = conditions.ts.pdb_path
            
            results.append(result)

            title = (
                f"[Record {record_idx}/{total_records}] "
                f"[Sample {sample_idx}/{run_count}] "
                f"selected_from={result['selected_from']}"
            )
            print("=" * 88)
            print(title)
            print("-" * 88)

            all_tasks = sorted(set(result["task_scores"].keys()) | set(result["task_acc"].keys()))
            for task_name in all_tasks:
                score_val = result["task_scores"].get(task_name)
                acc_val = result["task_acc"].get(task_name)
                score_str = f"{score_val:.4f}" if score_val is not None else "N/A"
                acc_str = f"{acc_val:.4f}" if acc_val is not None else "N/A"
                print(f"  - {task_name:<16} score={score_str:<8} acc={acc_str}")

            print(f"  - selected_acc      {result['selected_acc']:.4f}")
            print(f"  - length            {result['length']}")
            print(f"  - sequence          {result['sequence']}")
            print("=" * 88)

    output_path = pathlib.Path(args.output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print(json.dumps(results, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()