inference_generation_utils.py

import random
from typing import Any, Dict, List, Optional, Tuple

import torch

from .inference_imports import import_utils_ernie
from .inference_types import MODEL_SEQUENCE_LENGTH, TARGET_SEQUENCE_LENGTH
from .structure_parser import struct_to_feature_vector


def calculate_seq_similarity(seq1: str, seq2: str) -> float:
    if not seq1 or not seq2:
        return 0.0
    length = min(len(seq1), len(seq2))
    if length == 0:
        return 0.0
    same = sum(1 for i in range(length) if seq1[i] == seq2[i])
    return same / min(len(seq1), len(seq2))


def build_seq_cond_from_structure(secondary_structure: str, seq_len: int, device: torch.device) -> Optional[torch.Tensor]:
    if not secondary_structure:
        return None
    seq_cond_feature = struct_to_feature_vector(secondary_structure, seq_len)
    dim_cond = seq_cond_feature.size(1)
    cls_cond = torch.zeros(1, dim_cond, dtype=torch.float32)
    eos_cond = torch.zeros(1, dim_cond, dtype=torch.float32)
    return torch.cat([cls_cond, seq_cond_feature, eos_cond], dim=0).unsqueeze(0).to(device)


def build_seq_cond_from_icshape(icshape: str, device: torch.device) -> Optional[torch.Tensor]:
    if not icshape:
        return None
    values = [0.0] + list(map(float, icshape.split(","))) + [0.0]
    return torch.tensor(values, dtype=torch.float32, device=device).unsqueeze(0)


def parse_design_pos(design_pos_str: str) -> Optional[List[Tuple[int, int]]]:
    if not design_pos_str or not design_pos_str.strip():
        return None

    design_pos_str = design_pos_str.strip()
    if design_pos_str.startswith("(") and design_pos_str.endswith(")"):
        design_pos_str = design_pos_str[1:-1].strip()

    range_strs = [r.strip() for r in design_pos_str.split(",") if r.strip()]
    if not range_strs:
        raise ValueError("design_pos string is empty after parsing")

    ranges: List[Tuple[int, int]] = []
    for range_str in range_strs:
        if "-" not in range_str:
            raise ValueError(f"Each range must be in format 'start-end', got: {range_str}")

        parts = range_str.split("-")
        if len(parts) < 2 or (len(parts) == 2 and (not parts[0].strip() or not parts[1].strip())):
            raise ValueError(f"Invalid range format 'start-end': {range_str}")

        try:
            start = int(parts[0].strip())
            end = int(parts[1].strip())
            if start < 0 or end < start:
                raise ValueError(f"Invalid range: start={start}, end={end}. start must be >= 0 and end >= start")
            ranges.append((start, end))
        except (ValueError, IndexError) as e:
            raise ValueError(f"Failed to parse range '{range_str}': {e}")

    if not ranges:
        raise ValueError("No valid ranges found in design_pos string")

    return ranges


def build_initial_tokens(
    tokenizer,
    lens: int,
    template: Optional[str],
    seed: Optional[str],
    seed_mask_num: int,
    mask_ratio: float,
    device: torch.device,
    design_pos: Optional[List[Tuple[int, int]]] = None,
    max_lens = 1022
) -> Tuple[torch.Tensor, int]:
    seq_len = min(lens, 1022) + 2
    if template is not None:
        template = template[:1022]
        rna_ids = torch.from_numpy(tokenizer.tokenize(template, add_cls_eos=True)).unsqueeze(0).to(device)

        if design_pos is not None:
            mask_pos = torch.zeros(rna_ids.shape[1], dtype=torch.bool, device=device)
            for start_pos, end_pos in design_pos:
                token_start = start_pos + 1
                token_end = end_pos + 2
                token_end = min(token_end, rna_ids.shape[1] - 1)
                if token_start < token_end:
                    mask_pos[token_start:token_end] = True

            if mask_pos.any():
                mask_pos = mask_pos & (torch.rand(rna_ids.shape[1], device=device) < mask_ratio)
            rna_ids[0, mask_pos] = tokenizer.mask_id
        else:
            mask_pos = torch.rand(rna_ids.shape[1], device=device) < mask_ratio
            rna_ids[0, mask_pos] = tokenizer.mask_id
    elif seed is not None:
        rna_ids = torch.full((1, seq_len), tokenizer.mask_id, dtype=torch.long, device=device)
        seed_ids = torch.from_numpy(tokenizer.tokenize(seed, add_cls_eos=False)).unsqueeze(0).to(device)
        if seed_mask_num > 0:
            mask_pos = random.sample(range(seed_ids.shape[1]), min(seed_mask_num, seed_ids.shape[1]))
            seed_ids[0, mask_pos] = tokenizer.mask_id
        end = min(2 + seed_ids.shape[1], seq_len)
        valid = end - 2
        if valid > 0:
            rna_ids[0, 2:end] = seed_ids[0, :valid]
    else:
        rna_ids = torch.full((1, seq_len), tokenizer.mask_id, dtype=torch.long, device=device)

    rna_ids[0, 0] = tokenizer.cls_id
    rna_ids[0, -1] = tokenizer.eos_id
    return rna_ids, rna_ids.shape[1]


def collect_diffusion_func_cond(conditions) -> Optional[List[str]]:
    func_list: List[str] = []
    for value in [
        conditions.rfam.func_cond,
        conditions.proteins.func_cond,
        conditions.disease.func_cond,
        conditions.text.func_cond,
    ]:
        if value is None:
            continue
        if isinstance(value, list):
            func_list.extend(value)
        else:
            func_list.append(value)
    func_str = "; ".join(func_list)
    return [func_str] if func_str else None


def build_cas13_full_sequence(target_before: str, target_at_guide: str, target_after: str, guide_seq: str) -> str:
    target_sequence = f"{target_before}{target_at_guide}{target_after}"
    guide_extended = f"{'N' * 20}{guide_seq}{'N' * 20}"
    full_sequence = f"{target_sequence}{guide_extended}"
    if len(full_sequence) != MODEL_SEQUENCE_LENGTH:
        raise ValueError(f"CAS13: Unexpected full sequence length {len(full_sequence)}, expected {MODEL_SEQUENCE_LENGTH}")
    return full_sequence


def extract_guide_seq_from_full(full_sequence: str) -> str:
    full_sequence = full_sequence.upper().replace("T", "U")
    guide_start = TARGET_SEQUENCE_LENGTH
    return full_sequence[guide_start:]


def prepare_ss_extra_state(tasks: List[str], conditions, tokenizer, device: torch.device) -> Dict[str, Any]:
    if "ss" not in tasks:
        return {}
    if not conditions.ss.secondary_structure:
        raise ValueError("Task 'ss' requires secondary_structure condition")

    utils_ernie = import_utils_ernie()
    from .ss_evaluator import dotbracket_to_matrix
    from .rna_ss_utils import build_vid_to_base_channel, rnafm_logits_to_base_prob

    true_pair = dotbracket_to_matrix(conditions.ss.secondary_structure)
    true_pair_tensor = torch.tensor(true_pair, dtype=torch.float32, device=device).unsqueeze(0)
    return {
        "true_pair_tensor": true_pair_tensor,
        "vid_to_base_channel": build_vid_to_base_channel(tokenizer),
        "prepare_input_for_ernierna": utils_ernie.prepare_input_for_ernierna,
        "rnafm_logits_to_base_prob": rnafm_logits_to_base_prob,
    }


def normalize_score_output(raw_score: Any) -> Dict[str, float]:
    if isinstance(raw_score, dict):
        score = float(raw_score.get("score", 0.0))
        acc = float(raw_score.get("acc", 0.0))
        return {"score": score, "acc": acc}
    if isinstance(raw_score, tuple) and len(raw_score) >= 2:
        return {"score": float(raw_score[0]), "acc": float(raw_score[1])}
    return {"score": float(raw_score), "acc": 0.0}