deepchopper 1.3.0__cp310-abi3-macosx_11_0_arm64.whl

deepchopper/train.py

@@ -0,0 +1,109 @@
+import os
+from typing import TYPE_CHECKING, Any
+
+import hydra
+import lightning as L  # noqa: N812
+from lightning import Callback, LightningDataModule, LightningModule, Trainer
+from omegaconf import DictConfig
+
+from .utils import (
+    RankedLogger,
+    extras,
+    get_metric_value,
+    instantiate_callbacks,
+    instantiate_loggers,
+    log_hyperparameters,
+    task_wrapper,
+)
+
+if TYPE_CHECKING:
+    from lightning.pytorch.loggers import Logger
+
+import torch
+
+torch.set_float32_matmul_precision("high")
+
+log = RankedLogger(__name__, rank_zero_only=True)
+
+
+@task_wrapper
+def train(cfg: DictConfig) -> tuple[dict[str, Any], dict[str, Any]]:
+    """Trains the model. Can additionally evaluate on a testset, using best weights obtained during training.
+
+    This method is wrapped in optional @task_wrapper decorator, that controls the behavior during
+    failure. Useful for multiruns, saving info about the crash, etc.
+
+    :param cfg: A DictConfig configuration composed by Hydra.
+    :return: A tuple with metrics and dict with all instantiated objects.
+    """
+    # set seed for random number generators in pytorch, numpy and python.random
+    if cfg.get("seed"):
+        L.seed_everything(cfg.seed, workers=True)
+
+    log.info(f"Instantiating datamodule <{cfg.data._target_}>")
+    datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data)
+
+    log.info(f"Instantiating model <{cfg.model._target_}>")
+    model: LightningModule = hydra.utils.instantiate(cfg.model)
+
+    log.info("Instantiating callbacks...")
+    callbacks: list[Callback] = instantiate_callbacks(cfg.get("callbacks"))
+
+    log.info("Instantiating loggers...")
+    logger: list[Logger] = instantiate_loggers(cfg.get("logger"))
+
+    log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
+    trainer: Trainer = hydra.utils.instantiate(cfg.trainer, callbacks=callbacks, logger=logger)
+
+    object_dict = {
+        "cfg": cfg,
+        "datamodule": datamodule,
+        "model": model,
+        "callbacks": callbacks,
+        "logger": logger,
+        "trainer": trainer,
+    }
+
+    if logger:
+        log.info("Logging hyperparameters!")
+        log_hyperparameters(object_dict)
+
+    if cfg.get("train"):
+        log.info("Starting training!")
+        trainer.fit(model=model, datamodule=datamodule, ckpt_path=cfg.get("ckpt_path"))
+
+    train_metrics = trainer.callback_metrics
+
+    if cfg.get("test"):
+        log.info("Starting testing!")
+        ckpt_path = trainer.checkpoint_callback.best_model_path
+        if ckpt_path == "":
+            log.warning("Best ckpt not found! Using current weights for testing...")
+            ckpt_path = None
+        trainer.test(model=model, datamodule=datamodule, ckpt_path=ckpt_path)
+        log.info(f"Best ckpt path: {ckpt_path}")
+
+    test_metrics = trainer.callback_metrics
+
+    # merge train and test metrics
+    metric_dict = {**train_metrics, **test_metrics}
+
+    return metric_dict, object_dict
+
+
+@hydra.main(version_base="1.3", config_path=os.getenv("DC_CONFIG_PATH", "configs"), config_name="train.yaml")
+def main(cfg: DictConfig) -> float | None:
+    """Main entry point for training.
+
+    :param cfg: DictConfig configuration composed by Hydra.
+    :return: Optional[float] with optimized metric value.
+    """
+    # apply extra utilities
+    # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.)
+    extras(cfg)
+
+    # train the model
+    metric_dict, _ = train(cfg)
+
+    # safely retrieve metric value for hydra-based hyperparameter optimization
+    return get_metric_value(metric_dict=metric_dict, metric_name=cfg.get("optimized_metric"))

deepchopper/ui/__init__.py

@@ -0,0 +1 @@
+from .main import main

deepchopper/ui/main.py

@@ -0,0 +1,189 @@
+import multiprocessing
+from functools import partial
+from pathlib import Path
+
+import gradio as gr
+import lightning
+import torch
+from datasets import Dataset
+from torch.utils.data import DataLoader
+
+import deepchopper
+from deepchopper.deepchopper import default, remove_intervals_and_keep_left, smooth_label_region
+from deepchopper.models.llm import (
+    tokenize_and_align_labels_and_quals,
+)
+from deepchopper.utils import (
+    summary_predict,
+)
+
+
+def parse_fq_record(text: str):
+    """Parse a single FASTQ record into a dictionary."""
+    lines = text.strip().split("\n")
+    for i in range(0, len(lines), 4):
+        content = lines[i : i + 4]
+        record_id, seq, _, qual = content
+        assert len(seq) == len(qual)  # noqa: S101
+
+        input_qual = deepchopper.encode_qual(qual, default.QUAL_OFFSET)
+
+        yield {
+            "id": record_id,
+            "seq": seq,
+            "qual": torch.Tensor(input_qual),
+            "target": [0, 0],
+        }
+
+
+def load_dataset(text: str, tokenizer):
+    """Load dataset from text."""
+    dataset = Dataset.from_generator(parse_fq_record, gen_kwargs={"text": text}).with_format("torch")
+    tokenized_dataset = dataset.map(
+        partial(
+            tokenize_and_align_labels_and_quals,
+            tokenizer=tokenizer,
+            max_length=tokenizer.max_len_single_sentence,
+        ),
+        num_proc=multiprocessing.cpu_count(),  # type: ignore
+    ).remove_columns(["id", "seq", "qual", "target"])
+    return dataset, tokenized_dataset
+
+
+def predict(
+    text: str,
+    smooth_window_size: int = 21,
+    min_interval_size: int = 13,
+    approved_interval_number: int = 20,
+    max_process_intervals: int = 8,  # default is 4
+    batch_size: int = 1,
+    num_workers: int = 1,
+):
+    tokenizer = deepchopper.models.llm.load_tokenizer_from_hyena_model(model_name="hyenadna-small-32k-seqlen")
+    dataset, tokenized_dataset = load_dataset(text, tokenizer)
+
+    dataloader = DataLoader(tokenized_dataset, batch_size=batch_size, num_workers=num_workers, persistent_workers=True)
+    model = deepchopper.DeepChopper.from_pretrained("yangliz5/deepchopper")
+
+    accelerator = "gpu" if torch.cuda.is_available() else "cpu"
+    trainer = lightning.pytorch.trainer.Trainer(
+        accelerator=accelerator,
+        devices="auto",
+        deterministic=False,
+        logger=False,
+    )
+
+    predicts = trainer.predict(model=model, dataloaders=dataloader, return_predictions=True)
+
+    assert len(predicts) == 1  # noqa: S101
+
+    smooth_interval_json: list[dict[str, int]] = []
+    highlighted_text: list[tuple[str, str | None]] = []
+
+    for idx, preds in enumerate(predicts):
+        true_prediction, _true_label = summary_predict(predictions=preds[0], labels=preds[1])
+
+        _id = dataset[idx]["id"]
+        seq = dataset[idx]["seq"]
+
+        smooth_predict_targets = smooth_label_region(
+            true_prediction[0], smooth_window_size, min_interval_size, approved_interval_number
+        )
+
+        if not smooth_predict_targets or len(smooth_predict_targets) > max_process_intervals:
+            continue
+
+        # zip two consecutive elements
+        _selected_seqs, selected_intervals = remove_intervals_and_keep_left(seq, smooth_predict_targets)
+        total_intervals = sorted(selected_intervals + smooth_predict_targets)
+
+        smooth_interval_json.extend({"start": i[0], "end": i[1]} for i in smooth_predict_targets)
+
+        highlighted_text.extend(
+            (seq[interval[0] : interval[1]], "ada" if interval in smooth_predict_targets else None)
+            for interval in total_intervals
+        )
+    return smooth_interval_json, highlighted_text
+
+
+def process_input(text: str | None, file: str | None):
+    """Process the input and return the prediction."""
+    if not text and not file:
+        gr.Warning("Both text and file are empty")
+
+    if file:
+        MAX_LINES = 4
+        file_content = []
+        with Path(file).open() as f:
+            for idx, line in enumerate(f):
+                if idx >= MAX_LINES:
+                    break
+                file_content.append(line)
+        text = "".join(file_content)
+        return predict(text=text)
+
+    return predict(text=text)
+
+
+def create_gradio_app():
+    """Create a Gradio app for DeepChopper."""
+    example = (
+        "@1065:1135|393d635c-64f0-41ed-8531-12174d8efb28+f6a60069-1fcf-4049-8e7c-37523b4e273f\n"
+        "GCAGCTATGAATGCAAGGCCACAAGGTGGATGGAAGAGTTGTGGAACCAAAGAGCTGTCTTCCAGAGAAGATTTCGAGATAAGTCGCCCATCAGTGAACAAGATATTGTTGGTGGCATTTGATGAGAACGTTCCAAGATTATTGACAGATTAGTGAAAAGTAAGATTGAAATCATGACTGACCGTAAGTGGCAAGAAAGGGCTTTTGCCTTTGTAACCTTTGACGACCATGACTCCGTGGATAAGATTGTCATTCAGAATACCATACTGTGAATGGCCACATCTTTATTGTGAAGTTAGAAAAGCCCTGTCAAAGCAAGAGATGAATCAGTGCTTCTCCAGCCAAAGAGGTCGAAGTGGTTCTGGAAACTTTGGTGGTGGTCGTGGAGGTGGTTTCGGTGGGAATGACAACTCGGTCGTGGAGGAAACTTCAGTGGTCGTGGTGGCTTTGGTGGCAGCCGTGGTGGTGGTGGATATGGTGGCAGTGGGGATGGCTATAATGGATTTGGTAATGATGGAAGCAATTTGGAGGTGGTGGAAGCTACAATGATTTTGGGAATTACAACAATCAGTCTTCAAATTTTGGACCCCTAGGAGGAAATTTTGGTAGAAGCTCTGGCCCCATGGCGGTGGAGGCCAAATACTTTTGCAAACCACGAAACCAAGGTGGCTATGGCGGTCCAGCAGCAGCAGTAGCTATGGCAGTGGCAGAAGATTTTAATTAGGAAACAAAGCTTAGCAGGAGAGGAGAGCCAGAGAAGTGACAGGGAAGTACAGGTTACAACAGATTTGTGAACTCAGCCCAAGCACAGTGGTGGCAGGGCCTAGCTGCTACAAAGAAGACATGTTTTAGACAAATACTCATGTGTATGGGCAAAACTTGAGGACTGTATTTGTGACTAACTGTATAACAGGTTATTTTAGTTTCTGTTTGTGGAAAGTGTAAAGCATTCCAACAAAGGTTTTTAATGTAGATTTTTTTTTTTGCACCCCATGCTGTTGATTTGCTAAATGTAACAGTCTGATCGTGACGCTGAATAAATGTCTTTTTTAAAAAAAAAAAAAAGCTCCCTCCCATCCCCTGCTGCTAACTGATCCCATTATATCTAACCTGCCCCCCCATATCACCTGCTCCCGAGCTACCTAAGAACAGCTAAAAGAGCACACCCGCATGTAGCAAAATAGTGGGAAGATTATAGGTAGAGGCGACAAACCTACCGAGCCTGGTGATAGCTGGTTGTCCTAGATAGAATCTTAGTTCAACTTTAAATTTGCCCACAGAACCCTCTAAATCCCCTTGTAAATTTAACTGTTAGTCCAAAGAGGAACAGCTCTTTGGACACTAGGAAAAAACCTTGTAGAGAGTAAAAAATCAACACCCA\n"
+        "+\n"
+        ".0==?SSSSSSSSSSSH2216<868;SSSSSSSSSQQSRSIIHEDDESSSSSSJIKMGEKISSJJICCBDQ?;;8:;,**(&$'+501)\"#$()+%&&0<5+*/('%'))))'''$##\"\"\"\"%&--$\"\"\"('%)1L3*'')'#\"#&+*$&\"\"#*(&'''+,,<;9<BHGF//.LKORQSK<###%*-89<FSSSSE=BAFHFDB???3313NN?>=ANOSJDCADHGMOQSSD=7>BRRSPIEEEOQSSQ4->LIC7EE045///03IIJQSSSNGE6('.5??@A@=,,EGRSPKJ<==<556GFLLQRANSSSSSSSSG...*%%%(***(%'3@LOOSSSSM...7BCMMSSSSSSSSSSSSSSSDFIPSSSGGGGPOQLIHIL4103HMSILLNOSSSSSSSSSS22CBCGSHHHHSSSSSSSSD??@<<<:DDDSSSSSSSSSSA@6688OSSSSSROJJKLSNNNMSSSSQPOOSOOQSSSSSRRHIHISSRSSSSSSSSSSSJFF=??@SSQRK:424<444FFG///1S@@@ASNNNNPN:4JMDDLPSSSSSSBA?B?@@+'&'BD**8EDEFQPIMLE$$&',79CSJJPSGA+***DN;3-('&(;>6(()/-,,)%')1FRNNJ-:=>GC;&;CHNFFDCEEKJLFA22/27A.....HSQLHL))8<=?JSSSFGSKIHDDCCEFDAA@CFJKLNL>:9/1>>?OSLK@+HPSA;>>>K;;;;SSSSOQLPPMORSSSSSQSSSSSSS=:9**?D889SSRFFEDKJJJEEDKSSSNNOSSS.---,&*++SSSSQRSSSSQPGED<<89<@GJ999:SSKBBBAJHK=SSSJJKNMGHKKHQA<<>OPKFEAACDHJKMORB/)'((6**)15DA99;JSQSSS2())+J))EGMQOMMKJF>?<<AA620..D..,/112SOIIJSQFNEEEOMF?066=>@4,3;B>87FSSSSSSSSSSSSSSS<<::5658@AHMMSSRECC448/=<<>SSCB:5546;<??KF==;;FFEDFHKKJG):C>=>BJHINJFDPPPPPPPPPPPPPP%'*%$%+-%'(-22&&%('''&&&#\"\"%&'+0,,0;:1&\"\"%'(+++8'**(\"$$#&$'**//.3497$\"3CFHLOSSSSR:887:;;FSSRPRSSS4433$#$%&$$-056>@:;>=@?AHEFEC;*EKMSSRSRRDB>=AFRSSSSBSOOPSMDAABHH976951-9DHPQO/---?@ELSSQSRJHKKBKKLSSLINSOSSQSRIMSSSSSS>?MKIINSSGSSSSSSSQQMK544MJKKNKHGGLFFGBDB?EHIKGD?@DHPPIIF555)&(+,ADSSSSRQSSSQSS=9/0JJMSQSOSSO/97=B@=:>"
+    )
+
+    custom_css = """
+    .header { text-align: center; margin-bottom: 30px; }
+    .footer { text-align: center; margin-top: 30px; font-size: 0.8em; color: #666; }
+    """
+
+    with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
+        gr.HTML(
+            """
+            <div class="header">
+                <h1>🧬 DeepChopper: DNA Sequence Analysis</h1>
+                <p>Analyze DNA sequences and detect artificial sequences</p>
+            </div>
+            """
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                text_input = gr.Textbox(
+                    label="Input DNA Sequence", placeholder="Paste your DNA sequence here...", lines=10
+                )
+                file_input = gr.File(label="Or upload a FASTQ file")
+                submit_btn = gr.Button("Analyze", variant="primary")
+
+            with gr.Column(scale=1):
+                json_output = gr.JSON(label="Detected Artificial Regions")
+                highlighted_text = gr.HighlightedText(label="Highlighted Sequence")
+
+        submit_btn.click(fn=process_input, inputs=[text_input, file_input], outputs=[json_output, highlighted_text])
+
+        gr.Examples(
+            examples=[[example]],
+            inputs=[text_input],
+        )
+
+        gr.HTML(
+            """
+            <div class="footer">
+                <p>DeepChopper - Powered by AI for DNA sequence analysis</p>
+            </div>
+            """
+        )
+
+    return demo
+
+
+def main():
+    """Launch the Gradio app."""
+    app = create_gradio_app()
+    app.launch()
+
+
+if __name__ == "__main__":
+    main()

deepchopper/utils/__init__.py

@@ -0,0 +1,37 @@
+"""Utils."""
+
+from .instantiators import instantiate_callbacks, instantiate_loggers
+from .logging_utils import log_hyperparameters
+from .preprocess import load_kmer2id, load_safetensor, save_ndarray_to_safetensor
+from .print import (
+    alignment_predict,
+    highlight_target,
+    highlight_targets,
+    hightlight_predict,
+    hightlight_predicts,
+    summary_predict,
+)
+from .pylogger import RankedLogger
+from .rich_utils import print_config_tree
+from .utils import device, extras, get_metric_value, task_wrapper
+
+__all__ = [
+    "RankedLogger",
+    "alignment_predict",
+    "device",
+    "extras",
+    "get_metric_value",
+    "highlight_target",
+    "highlight_targets",
+    "hightlight_predict",
+    "hightlight_predicts",
+    "instantiate_callbacks",
+    "instantiate_loggers",
+    "load_kmer2id",
+    "load_safetensor",
+    "log_hyperparameters",
+    "print_config_tree",
+    "save_ndarray_to_safetensor",
+    "summary_predict",
+    "task_wrapper",
+]

deepchopper/utils/instantiators.py

@@ -0,0 +1,54 @@
+import hydra
+from lightning import Callback
+from lightning.pytorch.loggers import Logger
+from omegaconf import DictConfig
+
+from . import pylogger
+
+log = pylogger.RankedLogger(__name__, rank_zero_only=True)
+
+
+def instantiate_callbacks(callbacks_cfg: DictConfig) -> list[Callback]:
+    """Instantiates callbacks from config.
+
+    :param callbacks_cfg: A DictConfig object containing callback configurations.
+    :return: A list of instantiated callbacks.
+    """
+    callbacks: list[Callback] = []
+
+    if not callbacks_cfg:
+        log.warning("No callback configs found! Skipping..")
+        return callbacks
+
+    if not isinstance(callbacks_cfg, DictConfig):
+        raise TypeError("Callbacks config must be a DictConfig!")
+
+    for _, cb_conf in callbacks_cfg.items():
+        if isinstance(cb_conf, DictConfig) and "_target_" in cb_conf:
+            log.info(f"Instantiating callback <{cb_conf._target_}>")
+            callbacks.append(hydra.utils.instantiate(cb_conf))
+
+    return callbacks
+
+
+def instantiate_loggers(logger_cfg: DictConfig) -> list[Logger]:
+    """Instantiates loggers from config.
+
+    :param logger_cfg: A DictConfig object containing logger configurations.
+    :return: A list of instantiated loggers.
+    """
+    logger: list[Logger] = []
+
+    if not logger_cfg:
+        log.warning("No logger configs found! Skipping...")
+        return logger
+
+    if not isinstance(logger_cfg, DictConfig):
+        raise TypeError("Logger config must be a DictConfig!")
+
+    for _, lg_conf in logger_cfg.items():
+        if isinstance(lg_conf, DictConfig) and "_target_" in lg_conf:
+            log.info(f"Instantiating logger <{lg_conf._target_}>")
+            logger.append(hydra.utils.instantiate(lg_conf))
+
+    return logger

deepchopper/utils/logging_utils.py

@@ -0,0 +1,53 @@
+from typing import Any
+
+from lightning_utilities.core.rank_zero import rank_zero_only
+from omegaconf import OmegaConf
+
+from . import pylogger
+
+log = pylogger.RankedLogger(__name__, rank_zero_only=True)
+
+
+@rank_zero_only
+def log_hyperparameters(object_dict: dict[str, Any]) -> None:
+    """Controls which config parts are saved by Lightning loggers.
+
+    Additionally saves:
+        - Number of model parameters
+
+    :param object_dict: A dictionary containing the following objects:
+        - `"cfg"`: A DictConfig object containing the main config.
+        - `"model"`: The Lightning model.
+        - `"trainer"`: The Lightning trainer.
+    """
+    hparams = {}
+
+    cfg = OmegaConf.to_container(object_dict["cfg"])
+    model = object_dict["model"]
+    trainer = object_dict["trainer"]
+
+    if not trainer.logger:
+        log.warning("Logger not found! Skipping hyperparameter logging...")
+        return
+
+    hparams["model"] = cfg["model"]
+
+    # save number of model parameters
+    hparams["model/params/total"] = sum(p.numel() for p in model.parameters())
+    hparams["model/params/trainable"] = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    hparams["model/params/non_trainable"] = sum(p.numel() for p in model.parameters() if not p.requires_grad)
+
+    hparams["data"] = cfg["data"]
+    hparams["trainer"] = cfg["trainer"]
+
+    hparams["callbacks"] = cfg.get("callbacks")
+    hparams["extras"] = cfg.get("extras")
+
+    hparams["task_name"] = cfg.get("task_name")
+    hparams["tags"] = cfg.get("tags")
+    hparams["ckpt_path"] = cfg.get("ckpt_path")
+    hparams["seed"] = cfg.get("seed")
+
+    # send hparams to all loggers
+    for logger in trainer.loggers:
+        logger.log_hyperparams(hparams)

deepchopper/utils/preprocess.py

@@ -0,0 +1,62 @@
+from pathlib import Path
+
+import numpy as np
+import torch
+from safetensors.torch import load_file, save_file
+
+
+def load_safetensor(file_path: Path, device="cpu") -> dict[str, torch.Tensor]:
+    file_path = Path(file_path)
+    return load_file(file_path, device=device)
+
+
+def save_ndarray_to_safetensor(ndarray: dict[str, np.ndarray], file_path: Path):
+    """Save a dictionary of NumPy arrays as PyTorch tensors to a file.
+
+    Parameters:
+    - ndarray (dict[str, np.ndarray]): A dictionary where keys are strings and values are NumPy arrays.
+    - file_path (Path): The file path where the PyTorch tensors will be saved.
+
+    Returns:
+    - None
+
+    Raises:
+    - ValueError: If the input dictionary is empty.
+    - TypeError: If the input dictionary values are not NumPy arrays.
+    - FileNotFoundError: If the specified file path does not exist.
+    - OSError: If there is an issue with writing to the file.
+
+    This function converts each NumPy array in the input dictionary to a PyTorch tensor and saves the resulting dictionary
+    to a file specified by the file path. If any errors occur during the process, appropriate exceptions are raised.
+    """
+    tensor_dict = {key: torch.tensor(value) for key, value in ndarray.items()}
+    save_file(tensor_dict, file_path)
+
+
+def load_kmer2id(kmer2id_file: Path) -> dict:
+    """Load a dictionary mapping kmer strings to integer IDs from a file.
+
+    Parameters:
+    kmer2id_file (Path): A Path object pointing to the file containing the kmer to ID mapping.
+
+    Returns:
+    dict: A dictionary mapping kmer strings to integer IDs.
+
+    Raises:
+    FileNotFoundError: If the specified file does not exist.
+    ValueError: If the file format is incorrect or if the kmer and ID cannot be properly parsed.
+
+    Example:
+    If the kmer2id_file contains the following lines:
+    ATG 1
+    CAG 2
+    The function will return {'ATG': 1, 'CAG': 2}.
+    """
+    kmer2id_file = Path(kmer2id_file)
+
+    kmer2id = {}
+    with Path(kmer2id_file).open("r") as f:
+        for line in f:
+            kmer, idx = line.strip().split()
+            kmer2id[kmer] = int(idx)
+    return kmer2id

deepchopper/utils/print.py

@@ -0,0 +1,102 @@
+import typing
+
+import numpy as np
+import torch
+from rich.console import Console
+from rich.highlighter import RegexHighlighter
+from rich.text import Text
+from rich.theme import Theme
+
+from deepchopper.deepchopper import summary_predict as rust_summary_predict
+from deepchopper.models.llm import IGNORE_INDEX
+
+
+def hightlight_predicts(
+    seq: str,
+    targets: list[tuple[int, int]],
+    predicts: list[tuple[int, int]],
+    style: str = "bold magenta",
+    width: int = 80,
+):
+    """Highlight the predicted and labeled sequences."""
+    target_seq = Text(seq)
+    predict_seq = Text(seq)
+    console = Console()
+
+    for start, end in targets:
+        target_seq.stylize(style, start, end)
+
+    for start, end in predicts:
+        predict_seq.stylize(style, start, end)
+
+    front1 = "L:"
+    front2 = "P:"
+    for t1, t2 in zip(target_seq.wrap(console, width), predict_seq.wrap(console, width), strict=True):
+        console.print(front1, t1)
+        console.print(front2, t2)
+
+
+def highlight_targets(seq: str, targets: list[tuple[int, int]], style="bold magenta"):
+    """Highlight the target sequences."""
+    text = Text(seq)
+    console = Console()
+    for start, end in targets:
+        text.stylize(style, start, end)
+    console.print(text)
+
+
+def highlight_target(seq: str, start: int, end: int, style="bold magenta"):
+    """Highlight the target sequence."""
+    text = Text(seq)
+    console = Console()
+    text.stylize(style, start, end)
+    console.print(text)
+
+
+def hightlight_predict(seq: str, target_start: int, target_end: int, predict_start: int, predict_end: int):
+    """Highlight the predicted sequence."""
+    text = Text(seq)
+    console = Console()
+
+    text.stylize("#adb0b1", target_start, target_end)
+    text.stylize("bold magenta", predict_start, predict_end)
+
+    console.print(text)
+
+
+def summary_predict(predictions, labels):
+    """Summarize the predictions and labels."""
+    predictions = np.argmax(predictions, axis=2)
+    # Initialize lists to hold the filtered predictions and labels
+
+    if isinstance(predictions, torch.Tensor):
+        predictions = predictions.cpu().numpy()
+
+    if isinstance(labels, torch.Tensor):
+        labels = labels.cpu().numpy()
+
+    true_predictions, true_labels = rust_summary_predict(predictions, labels, IGNORE_INDEX)
+    return true_predictions, true_labels
+
+
+class LabelHighlighter(RegexHighlighter):
+    """Apply style to anything that looks like an email."""
+
+    base_style = "label."
+    highlights: typing.ClassVar = [r"(?P<label>1+)"]
+
+
+def alignment_predict(prediction, label):
+    """Print the alignment of the predicted and labeled sequences."""
+    import textwrap
+
+    prediction_str = "".join(str(x) for x in prediction)
+    label_str = "".join(str(x) for x in label)
+
+    front2 = "L:"
+    front1 = "P:"
+    theme = Theme({"label.label": "bold magenta"})
+    console = Console(highlighter=LabelHighlighter(), theme=theme)
+    for l1, l2 in zip(textwrap.wrap(prediction_str), textwrap.wrap(label_str), strict=True):
+        ss = f"{front1}{l1}\n{front2}{l2}"
+        console.print(ss)

deepchopper/utils/pylogger.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from lightning_utilities.core.rank_zero import rank_prefixed_message, rank_zero_only
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+
+class RankedLogger(logging.LoggerAdapter):
+    """A multi-GPU-friendly python command line logger."""
+
+    def __init__(
+        self,
+        name: str = __name__,
+        *,
+        rank_zero_only: bool = False,
+        extra: Mapping[str, object] | None = None,
+    ) -> None:
+        """Initializes a multi-GPU-friendly python command line logger that logs on all processes
+        with their rank prefixed in the log message.
+
+        :param name: The name of the logger. Default is ``__name__``.
+        :param rank_zero_only: Whether to force all logs to only occur on the rank zero process. Default is `False`.
+        :param extra: (Optional) A dict-like object which provides contextual information. See `logging.LoggerAdapter`.
+        """
+        logger = logging.getLogger(name)
+        super().__init__(logger=logger, extra=extra)
+        self.rank_zero_only = rank_zero_only
+
+    def log(self, level: int, msg: str, rank: int | None = None, *args, **kwargs) -> None:
+        """Delegate a log call to the underlying logger, after prefixing its message with the rank
+        of the process it's being logged from. If `'rank'` is provided, then the log will only
+        occur on that rank/process.
+
+        :param level: The level to log at. Look at `logging.__init__.py` for more information.
+        :param msg: The message to log.
+        :param rank: The rank to log at.
+        :param args: Additional args to pass to the underlying logging function.
+        :param kwargs: Any additional keyword args to pass to the underlying logging function.
+        """
+        if self.isEnabledFor(level):
+            msg, kwargs = self.process(msg, kwargs)
+            current_rank = getattr(rank_zero_only, "rank", None)
+            if current_rank is None:
+                raise RuntimeError("The `rank_zero_only.rank` needs to be set before use")
+            msg = rank_prefixed_message(msg, current_rank)
+            if self.rank_zero_only:
+                if current_rank == 0:
+                    self.logger.log(level, msg, *args, **kwargs)
+            else:
+                if rank is None:
+                    self.logger.log(level, msg, *args, **kwargs)
+                elif current_rank == rank:
+                    self.logger.log(level, msg, *args, **kwargs)