Cell2Sentence turns scRNA-seq data into text for LLM training.
import gdown

from src import utils

FILE_ID = "1PYUM59fKclw-aeN79oL5ghCkU4kn6XvN"


def download():
    """Download example scRNA-seq data from Google Drive."""
    download_url = f"https://drive.google.com/uc?id={FILE_ID}"
    output_filepath = utils.DATA_DIR / "dominguez_sample.h5ad"
    gdown.download(download_url, str(output_filepath), quiet=True)


if __name__ == "__main__":
    download()

retrieve_example_data.py
vandijklab/cell2sentence-ft
"""cell2sentence"""
__version__ = "0.0.1"
from src.csdata import CSData

src/__init__.py
vandijklab/cell2sentence-ft
#
# @author Rahul Dhodapkar
#
import zlib

import igraph as ig
import jellyfish
import numpy as np
import pandas as pd
from scipy import stats
from sklearn import model_selection


def zlib_ncd(s1, s2):
    """
    Return the zlib normalized compression distance between two strings
    """
    bs1 = bytes(s1, "utf-8")
    bs2 = bytes(s2, "utf-8")

    comp_cat = zlib.compress(bs1 + bs2)
    comp_bs1 = zlib.compress(bs1)
    comp_bs2 = zlib.compress(bs2)

    return (len(comp_cat) - min(len(comp_bs1), len(comp_bs2))) / max(
        len(comp_bs1), len(comp_bs2)
    )


class CSData:
    """
    Lightweight wrapper class to wrap cell2sentence results.
    """

    def __init__(self, vocab, sentences, cell_names, feature_names):
        self.vocab = vocab  # Ordered Dictionary: {gene_name: num_expressed_cells}
        self.sentences = sentences  # list of sentences
        self.cell_names = cell_names  # list of cell names
        self.feature_names = feature_names  # list of gene names
        self.distance_matrix = None
        self.distance_params = None
        self.knn_graph = None

    def create_distance_matrix(self, dist_type="jaro", prefix_len=20):
        """
        Calculate the distance matrix for the CSData object with the specified
        edit distance method. Currently supported: ("levenshtein").

        Distance caculated as d = 1 / (1 + x) where x is the similarity score.
        """
        if self.distance_matrix is not None and (
            self.distance_params["dist_type"] == dist_type
            and self.distance_params["prefix_len"] == prefix_len
        ):
            return self.distance_matrix

        dist_funcs = {
            "levenshtein": jellyfish.levenshtein_distance,
            "damerau_levenshtein": jellyfish.damerau_levenshtein_distance,
            "jaro": lambda x, y: 1 - jellyfish.jaro_similarity(x, y),  # NOQA
            "jaro_winkler": lambda x, y: 1
            - jellyfish.jaro_winkler_similarity(x, y),  # NOQA
            "zlib_ncd": zlib_ncd,
        }

        is_symmetric = {
            "levenshtein": True,
            "damerau_levenshtein": True,
            "jaro": True,
            "jaro_winkler": True,
            "zlib_ncd": False,
        }

        mat = np.zeros(shape=(len(self.sentences), len(self.sentences)))

        for i, s_i in enumerate(self.sentences):
            for j, s_j in enumerate(self.sentences):
                if j < i and is_symmetric[dist_type]:
                    mat[i, j] = mat[j, i]
                    continue

                mat[i, j] = dist_funcs[dist_type](s_i[:prefix_len], s_j[:prefix_len])

        self.distance_params = {"dist_type": dist_type, "prefix_len": prefix_len}
        self.distance_matrix = mat
        # reset KNN graph if previously computed on old distance
        self.knn_graph = None

        return self.distance_matrix

    def create_knn_graph(self, k=15):
        """
        Create KNN graph
        """
        if self.distance_matrix is None:
            raise RuntimeError(
                'cannot "build_knn_graph" without running "create_distance_matrix" first'
            )

        adj_matrix = 1 / (1 + self.distance_matrix)
        knn_mask = np.zeros(shape=adj_matrix.shape)

        for i in range(adj_matrix.shape[0]):
            for j in np.argsort(-adj_matrix[i])[:k]:
                knn_mask[i, j] = 1

        masked_adj_matrix = knn_mask * adj_matrix

        self.knn_graph = ig.Graph.Weighted_Adjacency(masked_adj_matrix).as_undirected()
        return self.knn_graph

    def create_rank_matrix(self):
        """
        Generates a per-cell rank matrix for use with matrix-based tools. Features with zero
        expression are zero, while remaining features are ranked according to distance from
        the end of the rank list.
        """
        full_rank_matrix = np.zeros((len(self.cell_names), len(self.feature_names)))

        for i, s in enumerate((self.sentences)):
            for rank_position, c in enumerate(s):
                full_rank_matrix[i, ord(c)] = len(s) - rank_position

        return full_rank_matrix

    def find_differential_features(self, ident_1, ident_2=None, min_pct=0.1):
        """
        Perform differential feature rank testing given a set of sentence indexes.
        If only one group is given, the remaining sentences are automatically used
        as the comparator group.
        """

        if ident_2 is None:
            ident_2 = list(set(range(len(self.sentences))).difference(set(ident_1)))

        full_rank_matrix = self.create_rank_matrix()
        feature_ixs_to_test = np.array(
            np.sum(full_rank_matrix > 0, axis=0) > min_pct * len(self.cell_names)
        ).nonzero()[0]

        stats_results = []
        for f in feature_ixs_to_test:
            wilcox_stat, pval = stats.ranksums(
                x=full_rank_matrix[ident_1, f], y=full_rank_matrix[ident_2, f]
            )
            stats_results.append(
                {
                    "feature": self.feature_names[f],
                    "w_stat": wilcox_stat,
                    "p_val": pval,
                    "mean_rank_group_1": np.mean(full_rank_matrix[ident_1, f]),
                    "mean_rank_group_2": np.mean(full_rank_matrix[ident_2, f]),
                }
            )
        return pd.DataFrame(stats_results)

    def get_rank_data_for_feature(self, feature_name, invert=False):
        """
        Return an array of ranks corresponding to the prescence of a gene within
        each cell sentence. If a gene is not present in a cell sentence, np.nan
        is returned for that cell.

        Note that this returns rank (1-indexed), not position within the underlying
        gene rank list string (0-indexed).
        """
        feature_code = -1
        for i, k in enumerate(self.vocab.keys()):
            if k == feature_name:
                feature_code = i
                break

        if feature_code == -1:
            raise ValueError(
                "invalid feature {} not found in vocabulary".format(feature_name)
            )
        feature_enc = chr(feature_code)

        rank_data_vec = np.full((len(self.cell_names)), np.nan)
        for i, s in enumerate(self.sentences):
            ft_loc = s.find(feature_enc)
            if invert:
                rank_data_vec[i] = len(s) - ft_loc if ft_loc != -1 else np.nan
            else:
                rank_data_vec[i] = ft_loc + 1 if ft_loc != -1 else np.nan

        return rank_data_vec

    def create_sentence_strings(self, delimiter=" "):
        """
        Convert internal sentence representation (arrays of ints) to traditional
        delimited character strings for integration with text-processing utilities.
        """
        if np.any([delimiter in x for x in self.feature_names]):
            raise ValueError(
                (
                    'feature names cannot contain sentence delimiter "{}", '
                    + "please re-format and try again"
                ).format(delimiter)
            )

        enc_map = list(self.vocab.keys())

        joined_sentences = []
        for s in self.sentences:
            joined_sentences.append(delimiter.join([enc_map[ord(x)] for x in s]))

        return np.array(joined_sentences, dtype=object)

    def create_sentence_lists(self):
        """
        Convert internal sentence representation (arrays of ints) to
        sentence lists compatible with gensim
        """
        enc_map = list(self.vocab.keys())

        joined_sentences = []
        for s in self.sentences:
            joined_sentences.append([enc_map[ord(x)] for x in s])

        return np.array(joined_sentences, dtype=object)

    def train_test_validation_split(
        self, train_pct=0.8, test_pct=0.1, val_pct=0.1, random_state=42
    ):
        """
        Create train, test, and validation splits of the data given the supplied
        percentages with a specified random state for reproducibility.

        Arguments:
            sentences: an numpy.ndarray of sentences to be split.
            train_pct: Default = 0.6. the percentage of samples to assign to the training set.
            test_pct: Default = 0.2. the percentage of samples to assign to the test set.
            val_pct: Default = 0.2. the percentage of samples to assign to the validation set.
        Return:
            (train_sentences, test_sentences, val_sentences) split from the
            originally supplied sentences array.
        """
        if train_pct + test_pct + val_pct != 1:
            raise ValueError(
                "train_pct = {} + test_pct = {} + val_pct = {} do not sum to 1.".format(
                    train_pct, test_pct, val_pct
                )
            )

        s_1 = test_pct
        s_2 = val_pct / (1 - test_pct)

        X = range(len(self.sentences))
        X_train, X_test = model_selection.train_test_split(
            X, test_size=s_1, random_state=random_state
        )

        X_train, X_val = model_selection.train_test_split(
            X_train, test_size=s_2, random_state=random_state
        )

        return (X_train, X_test, X_val)

src/csdata.py
vandijklab/cell2sentence-ft
import random


def construct_cell_type_template(cell_type):
    vowels = {"a", "e", "i", "o", "u", "A", "E", "I", "O", "U"}
    if cell_type[0] in vowels:
        cell_type = f"n {cell_type}"
    else:
        cell_type = f" {cell_type}"

    cell_type_templates = [
        "List the initial 100 genes associated with a{}: ",
        "Provide the top 100 genes for a{}: ",
        "Identify 100 genes corresponding to a{}: ",
        "Catalog the 100 primary genes for a{}: ",
        "Name the first set of 100 genes in a{}: ",
        "Outline 100 genes initially associated with a{}: ",
        "Specify the first 100 genes linked to a{}: ",
        "Show the leading 100 genes of a{}: ",
        "Enumerate the 100 starting genes for a{}: ",
        "List the 100 most highly expressed genes in a{}, ordered from highest to lowest expression: ",
        "Provide a ranking of the top 100 expressed genes in a{} by decreasing levels: ",
        "Detail the 100 genes with the greatest expression levels in a{}: ",
        "Identify the 100 genes in a{} with the highest expression, sorted in descending order: ",
        "Catalog the 100 most active genes in a{} cell, ordered by decreasing expression: ",
        "Name the 100 genes with peak expression levels in a{}, from highest to lowest: ",
        "Describe the 100 most abundantly expressed genes in a{}, in descending order: ",
        "Which 100 genes have the highest expression in a{}? List them in descending order: ",
        "Show the 100 most significantly expressed genes in a{}, organized by decreasing levels: ",
        "Enumerate the 100 genes with the strongest expression in a{}, ranked from highest to lowest: ",
        "List the top 100 genes by decreasing expression specifically for a{}: ",
        "Provide the first 100 highly expressed genes in a{}, sorted by expression level: ",
        "Identify the initial 100 genes with peak expression levels in a{}: ",
        "Catalog the top 100 expressed genes corresponding to a{}, from highest to lowest: ",
        "What are the first 100 genes with the greatest expression in a{}? ",
        "Name the top 100 genes sorted by decreasing expression for a{}: ",
        "Show the leading 100 genes by expression level associated with a{}: ",
        "Enumerate the first 100 genes by highest expression levels in a{}: ",
        "Specify the 100 primary genes in a{} ordered by decreasing expression: ",
        "Outline the initial set of 100 genes with high expression in a{}: ",
        "Detail the 100 starting genes for a{}, ranked by expression level: ",
        "Rank the first 100 genes by expression level found in a{}: ",
        "Describe the 100 most active genes initially associated with a{}: ",
        "Which are the first 100 genes in a{} sorted by decreasing expression? ",
        "Provide a ranking of the initial 100 genes by decreasing expression levels in a{}: ",
        "List 100 primary genes from highest to lowest expression in a{}: ",
        "Catalog the 100 initial genes for a{}, ordered by expression level from highest to lowest: ",
        "Identify the 100 leading genes in a{} based on expression levels: ",
        "Show the 100 primary genes for a{}, sorted by decreasing expression: ",
        "Enumerate the 100 most abundantly expressed starting genes in a{}: ",
    ]

    selected_template = random.choice(cell_type_templates)

    formatted_template = selected_template.format(cell_type)

    return formatted_template


def construct_prediction_template(genes):
    initial_prompt_templates = [
        "Identify the cell type most likely associated with these 100 highly expressed genes listed in descending order.",
        "Determine the probable cell type for the following 100 genes with the highest expression levels.",
        "Indicate the cell type typically linked to these 100 top-expressed genes.",
        "Specify the most likely cell type based on these 100 genes sorted by decreasing expression.",
        "Find the cell type that corresponds to these top 100 highly expressed genes.",
        "Point out the cell type these 100 genes with peak expression levels most commonly represent.",
        "Deduce the cell type likely to have these 100 highly expressed genes.",
        "Pinpoint the cell type that these 100 genes with the highest expression levels are most likely associated with.",
        "Ascertain the cell type from which these 100 highly expressed genes likely originate.",
        "Reveal the likely cell type linked to these 100 genes, listed by decreasing expression levels.",
        "Uncover the most probable cell type related to these 100 highly expressed genes.",
        "Indicate the cell type that would logically have these 100 top-expressed genes.",
        "Provide the likely cell type based on these 100 genes with high expression levels.",
        "Isolate the cell type commonly associated with these 100 top genes.",
        "Establish the cell type that these 100 genes with the highest expression levels are most likely from.",
        "Discern the likely cell type for these 100 genes sorted by expression level.",
        "Note the cell type typically associated with these 100 most expressed genes.",
        "Report the cell type most probably linked to these 100 genes with peak expression.",
        "Conclude the most likely cell type these 100 genes are associated with.",
        "State the probable cell type connected to these 100 top-expressed genes.",
        "What cell type is most likely represented by these top 100 highly expressed genes?",
        "Identify the probable cell type for these 100 genes with the highest expression levels.",
        "Which cell type is typically associated with these 100 most expressed genes?",
        "Can you deduce the cell type based on this list of 100 highly expressed genes?",
        "Given these 100 genes sorted by decreasing expression, what is the likely cell type?",
        "Based on these top 100 genes, which cell type are they most commonly found in?",
        "What type of cell is most likely to express these 100 genes in decreasing order of expression?",
        "What is the probable cell type these 100 highly expressed genes are associated with?",
        "From which cell type do these 100 most expressed genes likely originate?",
        "Determine the cell type likely associated with these 100 genes listed by decreasing expression.",
        "Given these 100 highly expressed genes, can you identify the likely cell type?",
        "Infer the cell type based on these 100 genes with the highest expression levels.",
        "Which cell type is likely to have these 100 genes with the highest expression?",
        "Could you specify the cell type most likely associated with these top 100 genes?",
        "What cell type would you associate with these 100 highly expressed genes?",
        "Can you tell the likely cell type for these 100 genes, sorted by decreasing expression?",
        "What is the likely cell type based on these 100 top expressed genes?",
        "Identify the cell type most commonly associated with these 100 genes.",
        "Based on these genes listed by decreasing expression, what cell type are they likely from?",
        "Given these 100 genes with high expression levels, what is the probable cell type?",
        "Which cell type is expected to have these 100 genes with the highest levels of expression?",
        "What is the most probable cell type based on these 100 genes with peak expression levels?",
        "What cell type would most likely have these 100 top expressed genes?",
        "Which cell type most probably corresponds to these 100 highly expressed genes?",
        "Could you determine the likely cell type based on these 100 most expressed genes?",
        "What type of cell would most likely contain these 100 genes with highest expression?",
        "Based on the list of 100 genes, what is the most likely corresponding cell type?",
        "Please identify the cell type that these 100 highly expressed genes are most likely linked to.",
        "Given these 100 genes ranked by expression, what would be the associated cell type?",
        "What would be the probable cell type for these 100 genes, listed by decreasing expression?",
        "Can you deduce the most likely cell type for these top 100 highly expressed genes?",
        "Identify the likely cell type these 100 genes with top expression could represent.",
        "Based on the following 100 genes, can you determine the cell type they are commonly found in?",
        "What is the likely originating cell type of these 100 top expressed genes?",
        "Specify the cell type most commonly linked with these 100 highly expressed genes.",
        "Which cell type would you expect to find these 100 genes with high expression levels?",
        "Indicate the probable cell type these 100 genes are commonly associated with.",
        "According to these 100 genes with highest expression, what cell type are they most likely from?",
        "Which cell type is these 100 genes with the highest expression levels most commonly found in?",
        "Could you point out the likely cell type linked with these 100 genes sorted by decreasing expression?",
    ]

    prediction_templates = [
        "This is the cell type corresponding to these genes: ",
        "These genes are most likely associated with the following cell type: ",
        "This is the probable cell type for these genes: ",
        "Based on these genes, the corresponding cell type is: ",
        "These genes suggest the cell type is most likely: ",
        "These genes are indicative of the cell type: ",
        "The associated cell type for these genes appears to be: ",
        "These genes typically correspond to: ",
        "The expected cell type based on these genes is: ",
        "These genes are commonly found in: ",
        "The cell type that these genes are most commonly linked with is: ",
        "Based on the expression levels, the cell type would likely be: ",
        "The genes provided are most commonly expressed in: ",
        "Given these genes, the likely corresponding cell type is: ",
        "The cell type these genes most likely originate from is: ",
        "These genes are most frequently associated with the cell type: ",
        "From these genes, it can be inferred that the cell type is: ",
        "The cell type best represented by these genes is: ",
    ]

    # build prompt

    selected_initial_template = random.choice(initial_prompt_templates)
    selected_prediction_template = random.choice(prediction_templates)

    formatted_template = (
        selected_initial_template + " " + genes + " " + selected_prediction_template
    )

    return formatted_template

src/prompts.py
vandijklab/cell2sentence-ft
#
# @author Rahul Dhodapkar <rahul.dhodapkar@yale.edu>
#

import os
import sys
from collections import OrderedDict
from pathlib import Path
from typing import List
from collections import Counter

import numpy as np
from scipy import sparse
from sklearn.utils import shuffle
from tqdm import tqdm

from src.csdata import CSData

DATA_DIR = Path("data/")
DATA_DIR.mkdir(exist_ok=True, parents=True)

BASE10_THRESHOLD = 3
SEED = 42


def generate_vocabulary(adata):
    """
    Create a vocabulary dictionary, where each key represents a single gene
    token and the value represents the number of non-zero cells in the provided
    count matrix.

    Arguments:
        adata: an AnnData object to generate cell sentences from. Expects that
               `obs` correspond to cells and `vars` correspond to genes.
    Return:
        a dictionary of gene vocabulary
    """
    if len(adata.var) > len(adata.obs):
        print(
            (
                "WARN: more variables ({}) than observations ({})... "
                + "did you mean to transpose the object (e.g. adata.T)?"
            ).format(len(adata.var), len(adata.obs)),
            file=sys.stderr,
        )

    vocabulary = OrderedDict()
    gene_sums = np.ravel(np.sum(adata.X > 0, axis=0))

    for i, name in enumerate(adata.var_names):
        vocabulary[name] = gene_sums[i]

    return vocabulary


def generate_sentences(adata, prefix_len=None, random_state=42):
    """
    Transform expression matrix to sentences. Sentences contain gene "words"
    denoting genes with non-zero expression. Genes are ordered from highest
    expression to lowest expression.

    Arguments:
        adata: an AnnData object to generate cell sentences from. Expects that
               `obs` correspond to cells and `vars` correspond to genes.
        random_state: sets the numpy random state for splitting ties
    Return:
        a `numpy.ndarray` of sentences, split by delimiter.
    """
    np.random.seed(random_state)

    if len(adata.var) > len(adata.obs):
        print(
            (
                "WARN: more variables ({}) than observations ({}), "
                + "did you mean to transpose the object (e.g. adata.T)?"
            ).format(len(adata.var), len(adata.obs)),
            file=sys.stderr,
        )

    mat = sparse.csr_matrix(adata.X)
    sentences = []
    for i in tqdm(range(mat.shape[0])):
        cols = mat.indices[mat.indptr[i] : mat.indptr[i + 1]]
        vals = mat.data[mat.indptr[i] : mat.indptr[i + 1]]

        cols, vals = shuffle(cols, vals)

        sentences.append(
            "".join([chr(x) for x in cols[np.argsort(-vals, kind="stable")]])
        )

    if prefix_len is not None:
        sentences = [s[:prefix_len] for s in sentences]

    return np.array(sentences, dtype=object)


def csdata_from_adata(adata, prefix_len=None, random_state=42):
    """
    Generate a CSData object from an AnnData object.

    Arguments:
        adata: an AnnData object to generate cell sentences from. Expects that
               `obs` correspond to cells and `vars` correspond to genes.
        prefix_len: consider only rank substrings of length prefix_len
        random_state: sets the numpy random state for splitting ties
    Return:
        a CSData object containing a vocabulary, sentences, and associated name data.
    """
    return CSData(
        vocab=generate_vocabulary(adata),
        sentences=generate_sentences(
            adata, prefix_len=prefix_len, random_state=random_state
        ),
        cell_names=adata.obs_names,
        feature_names=adata.var_names,
    )


def xlm_prepare_outpath(csdata, outpath, species_tag, params=None):
    """
    Write formatted data to the outpath file location, for direct processing
    by the XLM monolinguistic translation model. If creating an outpath for
    multiple species, use the same `outpath` with different `species_tag`
    values. They will not conflict so long as species_tags are appropriately
    assigned.

    Note that XLM requires a dictionary sorted in order of increasing
    frequency of occurence.

    Arguments:
        csdata: a CSData object from a single species to be written.
        outpath: directory to write files to. Will create this directory
                 if it does not already exist.
        species_tag: a short string to be used as the species name in XLM.
                     Fulfills functions analaglous to language tags such as
                     'en', 'es', or 'zh'.
        delimiter: default = ' '. A token delimter for the generated sentences.
        params: a parameter object passed to train_test_validation_split:
    Return:
        None
    """

    if params is None:
        params = {}

    sentence_strings = csdata.create_sentence_strings(delimiter=" ")
    train, test, val = csdata.train_test_validation_split(**params)

    train_sentences = sentence_strings[train]
    test_sentences = sentence_strings[test]
    val_sentences = sentence_strings[val]

    os.makedirs(outpath, exist_ok=True)
    np.save(
        os.path.join(outpath, "train_partition_indices.npy"),
        np.array(train, dtype=np.int64),
    )
    np.save(
        os.path.join(outpath, "valid_partition_indices.npy"),
        np.array(val, dtype=np.int64),
    )
    np.save(
        os.path.join(outpath, "test_partition_indices.npy"),
        np.array(test, dtype=np.int64),
    )

    print("INFO: Writing Vocabulary File", file=sys.stderr)
    fn = "{}/vocab_{}.txt".format(outpath, species_tag)
    with open(fn, "w") as f:
        for k in tqdm(sorted(csdata.vocab, key=csdata.vocab.get, reverse=True)):
            if csdata.vocab[k] == 0:
                continue
            print("{} {}".format(k, csdata.vocab[k]), file=f)

    print("INFO: Writing Training Sentences", file=sys.stderr)
    fn = "{}/train_{}.txt".format(outpath, species_tag)
    with open(fn, "w") as f:
        for l in tqdm(train_sentences):
            print(l, file=f)

    print("INFO: Writing Training Cell Barcodes", file=sys.stderr)
    fn = "{}/train_barcodes_{}.txt".format(outpath, species_tag)
    with open(fn, "w") as f:
        for l in tqdm(csdata.cell_names[train]):
            print(l, file=f)

    print("INFO: Writing Testing Sentences", file=sys.stderr)
    fn = "{}/test_{}.txt".format(outpath, species_tag)
    with open(fn, "w") as f:
        for l in tqdm(test_sentences):
            print(l, file=f)

    print("INFO: Writing Testing Cell Barcodes", file=sys.stderr)
    fn = "{}/train_barcodes_{}.txt".format(outpath, species_tag)
    with open(fn, "w") as f:
        for l in tqdm(csdata.cell_names[test]):
            print(l, file=f)

    print("INFO: Writing Validation Sentences", file=sys.stderr)
    fn = "{}/valid_{}.txt".format(outpath, species_tag)
    with open(fn, "w") as f:
        for l in tqdm(val_sentences):
            print(l, file=f)

    print("INFO: Writing Validation Cell Barcodes", file=sys.stderr)
    fn = "{}/valid_barcodes_{}.txt".format(outpath, species_tag)
    with open(fn, "w") as f:
        for l in tqdm(csdata.cell_names[val]):
            print(l, file=f)


def post_process_generated_cell_sentences(
    cell_sentence: str,
    global_dictionary: List,
    replace_nonsense_string: str = "NOT_A_GENE",
):
    """
    Post-processing function for generated cell sentences. Nonsense genes are replaced with 
    some string, e.g. 'NOT_A_GENE', so that ranks are not changed in generated output.

    Current assumptions in this function:
        - We replace nonsense genes with some string, e.g. 'NOT_A_GENE', so that ranks are not
            changed in generated output.

    Steps:
        1. Replace any nonsense genes with a specified token, e.g. 'NOT_A_GENE'
        2. Average the ranks of duplicated genes in generated sentence

    Arguments:
        cell_sentence:              generated cell sentence string
        global_dictionary:          list of global gene vocabulary (all uppercase)
        replace_nonsense_string:    string which will replace nonsense genes in generated output

    Returns:
        post_processed_sentence:    generated cell sentence after post processing steps
        num_nonsense_genes:         number of genes replaced with defined nonsense token
    """
    generated_gene_names = cell_sentence.split(" ")
    generated_gene_names = [generated_gene.upper() for generated_gene in generated_gene_names]

    # --- Replace nonsense genes ---#
    generated_gene_names = [
        gene_name if gene_name in global_dictionary else replace_nonsense_string
        for gene_name in generated_gene_names
    ]
    num_genes_replaced = generated_gene_names.count(replace_nonsense_string)

    # --- Average ranks ---#
    gene_name_to_occurrences = Counter(
        generated_gene_names
    )  # get mapping of gene name --> number of occurrences
    post_processed_sentence = generated_gene_names.copy()  # copy of generated gene list

    for gene_name in gene_name_to_occurrences:
        if (
            gene_name_to_occurrences[gene_name] > 1
            and gene_name != replace_nonsense_string
        ):
            # Find positions of all occurrences of duplicated generated gene in list
            # Note: using post_processed_sentence here; since duplicates are being removed, list will be
            #   getting shorter. Getting indices in original list will no longer be accurate positions
            occurrence_positions = [
                idx
                for idx, elem in enumerate(post_processed_sentence)
                if elem == gene_name
            ]
            average_position = int(
                sum(occurrence_positions) / len(occurrence_positions)
            )

            # Remove occurrences
            post_processed_sentence = [
                elem for elem in post_processed_sentence if elem != gene_name
            ]
            # Reinsert gene_name at average position
            post_processed_sentence.insert(
                average_position, gene_name
            )

    return post_processed_sentence, num_genes_replaced


def convert_cell_sentence_back_to_expression_vector(
    cell_sentence: List, global_dictionary: List, slope: float, intercept: float
):
    """
    Function to convert

    Current assumptions in this function:
        - We replace nonsense genes with some string, e.g. 'NOT_A_GENE', so that ranks are not
            changed in generated output.

    Steps:
        1. Replace any nonsense genes with a specified token, e.g. 'nan'
        2. Average the ranks of duplicated genes in generated sentence

    Arguments:
        cell_sentence:              generated cell sentence list, e.g. ['GENE1', 'GENE2']
        global_dictionary:          list of global gene vocabulary
        slope:                      slope value to use in inverse rank->expression transformation
        intercept:                  intercept value to use in inverse rank->expression transformation

    Returns:
        expression_vector:          expression vector for generated cell
    """
    expression_vector = np.zeros(len(global_dictionary), dtype=np.float32)
    for rank, gene_name in enumerate(cell_sentence):
        if gene_name in global_dictionary:
            log_rank = np.log10(1 + rank).item()
            gene_expr_val = intercept + (slope * log_rank)
            gene_idx_in_vector = global_dictionary.index(gene_name)
            expression_vector[gene_idx_in_vector] = gene_expr_val

    return expression_vector

src/utils.py
vandijklab/cell2sentence-ft
import json
import logging
import math
import os
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from itertools import chain
from random import choice
from typing import Optional

import numpy as np
import torch
import transformers
import wandb
from datasets import concatenate_datasets, load_from_disk
from torch import _dynamo as dynamo
from torch.nn import CrossEntropyLoss
from torch.utils import cpp_extension
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
)

from src.prompts import construct_cell_type_template, construct_prediction_template

LOG_LEVELS = {
    "debug": logging.DEBUG,
    "info": logging.INFO,
    "warning": logging.WARNING,
    "error": logging.ERROR,
    "critical": logging.CRITICAL,
}

if "LOCAL_RANK" in os.environ:
    LOCAL_RANK = int(os.environ["LOCAL_RANK"])
else:
    LOCAL_RANK = 0

logger = logging.getLogger(__name__)


@dataclass
class CustomTrainingArguments:
    model_name: str = field(
        default="gpt2", metadata={"help": "Hugging Face model name."}
    )
    seed: int = field(
        default=42, metadata={"help": "Seed for shuffling training dataset."}
    )
    data_seed: Optional[int] = field(
        default=None,
        metadata={"help": "Data seed for Hugging Face's trainer pipeline."},
    )
    set_torch_seed_manually: bool = field(
        default=False, metadata={"help": "Seed for PyTorch."}
    )
    torch_cuda_seed: int = field(
        default=42, metadata={"help": "Seed for PyTorch CUDA."}
    )
    eval_dataset_size: int = field(
        default=1000,
        metadata={"help": "Number of samples to use from evaluation dataset."},
    )
    evaluation_strategy: str = field(
        default="no",
        metadata={"help": "Whether to evaluate on steps, epochs, or none."},
    )
    eval_steps: int = field(
        default=100,
        metadata={
            "help": "If evaluation_strategy is set to 'steps', will evaluate every number of steps here."
        },
    )
    eval_accumulation_steps: int = field(
        default=5,
        metadata={"help": "Number of evaluation steps before offloading to CPU."},
    )
    output_dir: str = field(
        default="data/model/",
        metadata={"help": "Output directory for training runs."},
    )
    overwrite_output_dir: bool = field(
        default=False,
        metadata={"help": "Whether to overwrite output directory if nonempty."},
    )
    save_strategy: str = field(
        default="steps",
        metadata={
            "help": "Whether to save model checkpoints on steps, epochs, or none."
        },
    )
    save_steps: int = field(
        default=500,
        metadata={
            "help": "If save_strategy is set to 'steps', will save model checkpoint every number of steps here."
        },
    )
    save_total_limit: int = field(
        default=100,
        metadata={
            "help": "Maximum number of model checkpoints saved in output directory."
            " Will overwrite earlier checkpoints if limit is exceeded."
        },
    )
    per_device_train_batch_size: int = field(
        default=16, metadata={"help": "Per device batch size used during training."}
    )
    per_device_eval_batch_size: int = field(
        default=16, metadata={"help": "Per device batch size used during evaluation."}
    )
    num_train_epochs: int = field(
        default=5, metadata={"help": "Number of training epochs."}
    )
    wandb_project_name: str = field(
        default="cell2sentence", metadata={"help": "Wandb project name to save to."}
    )
    checkpoint: Optional[str] = field(
        default=None,
        metadata={"help": "Path to model checkpoint if resuming training."},
    )
    torch_compile: bool = field(
        default=False, metadata={"help": "Whether to use torch compile."}
    )
    torchdynamo: Optional[str] = field(
        default=None, metadata={"help": "Backend compiler for torch dynamo."}
    )
    torch_compile_backend: Optional[str] = field(
        default=None, metadata={"help": "Backend compiler for torch compile."}
    )
    dynamo_cache_size_limit: int = field(
        default=64, metadata={"help": "Number of graphs to cache for torch compile."}
    )
    dynamo_verbose: bool = field(
        default=False, metadata={"help": "Make dynamo config set to verbose."}
    )
    fp16: bool = field(default=False, metadata={"help": "Whether to use fp16."})
    ddp_backend: str = field(
        default="nccl", metadata={"help": "Backend for distributed data parallelism."}
    )
    dataloader_num_workers: int = field(
        default=0, metadata={"help": "Number of workers to use for dataloader."}
    )
    gradient_checkpointing: bool = field(
        default=False,
        metadata={
            "help": "Whether to checkpoint gradients during training. Improves GPU memory consumption."
        },
    )
    gradient_accumulation_steps: int = field(
        default=1,
        metadata={
            "help": "Number of forward passes before backpropagation during training. Per device."
        },
    )
    logging_steps: int = field(
        default=100,
        metadata={
            "help": "Number of training steps before logging, where steps is the number of gradient unpdates."
        },
    )
    data_dir: str = field(
        default="data/cell_sentences_hf/",
        metadata={"help": "Input directory for dataset."},
    )
    wandb_logging: bool = field(
        default=False, metadata={"help": "Whether to log to wandb."}
    )
    wandb_run_base_name: str = field(
        default="pbmc_finetune",
        metadata={"help": "Base name for wandb run. Start time will be appended."},
    )
    log_level: str = field(default="debug", metadata={"help": "Log level to use."})
    optim: str = field(
        default="adamw_torch",
        metadata={
            "help": "Optimizer to use. See Hugging Face options in TrainerArguments."
        },
    )
    deepspeed: str = field(
        default=None,
        metadata={"help": "Whether to use deepspeed for distributed training."},
    )


def main():
    if LOCAL_RANK == 0:
        logger.info(f"\nCUDA HOME: {cpp_extension.CUDA_HOME}")
        logger.info(f"\nTORCH CUDA VERSION: {torch.version.cuda}")

    logger.info(f"\nLOCAL RANK: {LOCAL_RANK}")

    assert torch.cuda.is_available(), "CUDA unavailable"
    device = torch.device("cuda")

    parser = HfArgumentParser((CustomTrainingArguments,))
    training_args = parser.parse_args_into_dataclasses(return_remaining_strings=True)[0]
    training_args_dict = asdict(training_args)

    if LOCAL_RANK == 0:
        logger.info(json.dumps(training_args_dict, indent=2))

    log_level = LOG_LEVELS[training_args.log_level]
    logger.setLevel(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    if training_args.set_torch_seed_manually:
        torch.cuda.manual_seed(training_args.torch_cuda_seed)
        logger.info(
            f"\nSET TORCH CUDA SEED MANUALLY. SEED VALUE: {training_args.torch_cuda_seed}"
        )

    if training_args.torch_compile:
        dynamo.config.cache_size_limit = training_args.dynamo_cache_size_limit
        dynamo.config.verbose = training_args.dynamo_verbose
    else:
        training_args.torchdynamo = None
        training_args.torch_compile_backend = None

    # load arrow datasets from dir
    dataset = load_from_disk(training_args.data_dir)
    train_dataset, val_dataset = dataset["train"].shuffle(
        seed=training_args.seed
    ), dataset["valid"].shuffle(seed=training_args.seed)

    if LOCAL_RANK == 0:
        logger.info(f"\nLENGTH OF TRAIN DATASET: {len(train_dataset)}")
        logger.info(train_dataset)

        logger.info(f"\nLENGTH OF EVAL DATASET: {len(val_dataset)}")
        logger.info(val_dataset)

    # Instantiate tokenizer and model for finetuning
    tokenizer = AutoTokenizer.from_pretrained(training_args.model_name)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(training_args.model_name).to(device)

    # Get current time and initialize wandb
    now = datetime.now()
    now = datetime.strftime(now, "%Y-%m-%d_%H-%M-%S")
    run_name = f"{training_args.model_name}-{training_args.wandb_run_base_name}-{now}"

    if training_args.wandb_logging:
        if LOCAL_RANK == 0:
            wandb.init(project=training_args.wandb_project_name, name=run_name)
            wandb.watch(model, log="all", log_freq=10)

    def preprocess_function(examples):
        text_column = "cell_type"
        label_column = "input_ids"
        max_length = 1024

        batch_size = len(examples[text_column])
        inputs = []
        targets = []
        for i in range(batch_size):
            prompt_type = choice([0, 1, 2])
            if prompt_type == 0:
                input = construct_cell_type_template(examples["cell_type"][i])
                target = " ".join(examples["input_ids"][i].split(" ")[:100])
            elif prompt_type == 1:
                input = construct_cell_type_template("PBMC")
                target = " ".join(examples["input_ids"][i].split(" ")[:100])
            else:
                input = construct_prediction_template(
                    " ".join(examples["input_ids"][i].split(" ")[:100])
                )
                target = examples["cell_type"][i]

            inputs.append(input)
            targets.append(target)

        model_inputs = tokenizer(inputs)
        labels = tokenizer(targets)

        for i in range(batch_size):
            sample_input_ids = model_inputs["input_ids"][i]
            label_input_ids = labels["input_ids"][i] + [tokenizer.eos_token_id]
            model_inputs["input_ids"][i] = sample_input_ids + label_input_ids
            labels["input_ids"][i] = [-100] * len(sample_input_ids) + label_input_ids
            model_inputs["attention_mask"][i] = [1] * len(model_inputs["input_ids"][i])

        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    train_dataset = train_dataset.map(preprocess_function, batched=True)
    val_dataset = val_dataset.map(preprocess_function, batched=True)
    if training_args.evaluation_strategy != "no":
        val_dataset = val_dataset.select(range(training_args.eval_dataset_size))

    # Collate function for training.
    def data_collator(examples):
        max_length = 0
        for i in range(len(examples)):
            input_length = len(examples[i]["input_ids"])
            max_length = max(max_length, input_length)

        batch_input_ids = []
        batch_attention_mask = []
        batch_labels = []

        for i in range(len(examples)):
            sample_input_ids = examples[i]["input_ids"]
            label_input_ids = examples[i]["labels"]
            attention_mask = examples[i]["attention_mask"]

            final_input_ids = [tokenizer.pad_token_id] * (
                max_length - len(sample_input_ids)
            ) + sample_input_ids
            final_attention_mask = [0] * (
                max_length - len(sample_input_ids)
            ) + attention_mask
            final_label_input_ids = [-100] * (
                max_length - len(sample_input_ids)
            ) + label_input_ids

            batch_input_ids.append(final_input_ids)
            batch_attention_mask.append(final_attention_mask)
            batch_labels.append(final_label_input_ids)

        return {
            "input_ids": torch.tensor(batch_input_ids),
            "attention_mask": torch.tensor(batch_attention_mask),
            "labels": torch.tensor(batch_labels),
        }

    # Configure Trainer and start training
    output_dir = training_args.output_dir + f"/{run_name}"

    train_args = TrainingArguments(
        output_dir=output_dir,
        overwrite_output_dir=training_args.overwrite_output_dir,
        seed=training_args.seed,
        data_seed=training_args.data_seed,
        per_device_train_batch_size=training_args.per_device_train_batch_size,
        per_device_eval_batch_size=training_args.per_device_eval_batch_size,
        evaluation_strategy=training_args.evaluation_strategy,
        eval_steps=training_args.eval_steps,
        eval_accumulation_steps=training_args.eval_accumulation_steps,
        num_train_epochs=training_args.num_train_epochs,
        report_to="wandb",
        torch_compile=training_args.torch_compile,
        torchdynamo=training_args.torchdynamo,
        torch_compile_backend=training_args.torch_compile_backend,
        fp16=training_args.fp16,
        ddp_backend=training_args.ddp_backend,
        dataloader_num_workers=training_args.dataloader_num_workers,
        gradient_checkpointing=training_args.gradient_checkpointing,
        gradient_accumulation_steps=training_args.gradient_accumulation_steps,
        logging_steps=training_args.logging_steps,
        save_strategy=training_args.save_strategy,
        save_steps=training_args.save_steps,
        save_total_limit=training_args.save_total_limit,
        optim=training_args.optim,
        deepspeed=training_args.deepspeed,
    )

    trainer = Trainer(
        model=model,
        args=train_args,
        data_collator=data_collator,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
    )
    if LOCAL_RANK == 0:
        logger.info(f"\nDEEPSPEED ENABLED: {trainer.is_deepspeed_enabled}")
        logger.info(f"\nFINAL TRAINING ARGUMENTS: {trainer.args}")
    train_result = trainer.train(resume_from_checkpoint=training_args.checkpoint)
    trainer.save_model()
    trainer.log_metrics("train", train_result.metrics)
    trainer.save_metrics("train", train_result.metrics)
    trainer.save_state()


if __name__ == "__main__":
    main()

train.py
vandijklab/cell2sentence-ft
import os
import argparse
from pathlib import Path

import anndata
import numpy as np
import pandas as pd
import plotnine as pn
import scanpy as sc
import sklearn.linear_model as lm
from datasets import Dataset, load_dataset, concatenate_datasets
from scipy.stats import pearsonr, spearmanr
from sklearn.metrics import r2_score
from sklearn.utils import shuffle
from tqdm import tqdm

from src import utils

ROW_SUM = 10000


def normalize_and_rank_transform(data_matrix_X, normalize=True):
    """
    Helper function which accepts a data matrix, optionally row-normalizes it,
    and calculated a rank transformation of the data.

    Args:
        data_matrix_X:  numpy matrix of shape [num_cells, num_genes]
        normalize:      boolean flag for whether to normalize data

    Returns:
        data_matrix_X:  normalized data matrix
        rank_matrix_X:  matrix of rank values for each cell, shame shape as data_matrix_X
    """
    if normalize:
        normalized_data_matrix_X = (
            np.diag(ROW_SUM / np.ravel(np.sum(data_matrix_X, axis=1))) @ data_matrix_X
        )
        data_matrix_X = np.asarray(normalized_data_matrix_X)

    rank_matrix_X = np.zeros(shape=data_matrix_X.shape)
    for i in tqdm(range(data_matrix_X.shape[0])):
        cols = np.ravel(range(data_matrix_X.shape[1]))
        vals = np.ravel(data_matrix_X[i, :])
        cols, vals = shuffle(cols, vals)
        ranks = cols[np.argsort(-vals, kind="stable")]
        for j in range(len(ranks)):
            rank_matrix_X[i, ranks[j]] = j

    return data_matrix_X, rank_matrix_X


def evaluate_transformation(df, plotting_sample_size=10000):
    """
    Helper function which takes as input a pandas DataFrame of expression values and
    ranks, and fits a linear regression model to predict back expression value from
    log rank.

    Plots are created to show the relationship between log rank and log expression,
    as well as the performance of expression reconstruction by the linear model.
    Metrics for expression reconstruction, as well as the parameters of the linear
    model are saved in a CSV file.

    Args:
        df:                     pandas DataFrame with keys: 'preprocessed_transcript_count,
                                    'preprocessed_rank', 'log_preprocessed_transcript_count',
                                    and 'log_preprocessed_rank'
        plotting_sample_size:   how many values to sample for plotting
    """
    eval_output_dir = utils.DATA_DIR / "eval"
    eval_output_dir.mkdir(exist_ok=True, parents=True)

    # (1) Fit linear regression between log rank (x-axis) and log expression (y-axis)
    x_axis_name = "log_preprocessed_rank"
    y_axis_name = "log_preprocessed_transcript_count"
    x = np.array(df.loc[df[x_axis_name] < utils.BASE10_THRESHOLD, x_axis_name]).reshape(
        -1, 1
    )
    y = df.loc[df[x_axis_name] < utils.BASE10_THRESHOLD, y_axis_name]

    reg = lm.LinearRegression().fit(x, y)

    # Plot relationship
    plot = (
        pn.ggplot(
            df.sample(plotting_sample_size),
            pn.aes(x="log_preprocessed_rank", y="log_preprocessed_transcript_count"),
        )
        + pn.geom_abline(slope=reg.coef_, intercept=reg.intercept_, color="red")
        + pn.geom_point(color="blue", size=0.5)
        + pn.labs(
            x="Gene Log Rank",
            y="Gene Log Expression",
            title="Log Rank vs Log Expression",
        )
    )
    plot.save(os.path.join(eval_output_dir, "plot_log_rank_vs_log_expr.png"), dpi=300)

    # (2) Reconstruct expression from log rank, calculate reconstruction performance metrics
    rank_reconstructed_X = reg.predict(
        np.array(df["log_preprocessed_rank"]).reshape(-1, 1)
    )

    r_squared_score = r2_score(
        np.asarray(df["log_preprocessed_transcript_count"]),
        np.asarray(rank_reconstructed_X),
    )
    pearson_r_score = pearsonr(
        np.asarray(df["log_preprocessed_transcript_count"]),
        np.asarray(rank_reconstructed_X),
    )
    spearman_r_score = spearmanr(
        np.asarray(df["log_preprocessed_transcript_count"]),
        np.asarray(rank_reconstructed_X),
    )

    reconstructed_expr_values_df = pd.DataFrame(
        {
            "Ground Truth Expression": df["log_preprocessed_transcript_count"],
            "Reconstructed Expression from Log Rank": rank_reconstructed_X,
        }
    )
    plot = (
        pn.ggplot(
            reconstructed_expr_values_df.sample(plotting_sample_size),
            pn.aes(
                x="Ground Truth Expression", y="Reconstructed Expression from Log Rank"
            ),
        )
        + pn.geom_point(color="blue", size=0.5)
        + pn.geom_abline(slope=1, intercept=0, color="red")
        + pn.labs(
            x="Ground Truth Expression",
            y="Reconstructed Expression from Log Rank",
            title="Ground Truth Expression vs Reconstruction from Rank",
        )
    )
    plot.save(
        os.path.join(
            eval_output_dir, "plot_gt_expr_vs_reconstructed_expr_from_rank.png"
        ),
        dpi=300,
    )

    # 3. Create results dataframe and return
    metrics_df = pd.DataFrame(
        {
            "threshold": [utils.BASE10_THRESHOLD],
            "slope": [reg.coef_.item()],
            "intercept": [reg.intercept_.item()],
            "R^2": [r_squared_score.item()],
            "Pearson_R_statistic": [pearson_r_score.statistic.item()],
            "Pearson_R_p_value": [pearson_r_score.pvalue.item()],
            "Spearman_R_statistic": [spearman_r_score.statistic.item()],
            "Spearman_R_p_value": [spearman_r_score.pvalue.item()],
        }
    )
    metrics_df.to_csv(
        os.path.join(eval_output_dir, "transformation_metrics_and_parameters.csv")
    )


def main(data_filepath: Path, output_dir: Path):
    """Apply preprocessing steps and transform to cell sentences.

    Preprocessing follows https://scanpy-tutorials.readthedocs.io/en/latest/pbmc3k.html.
    """
    print(f"Loading data from {data_filepath}.")
    adata = anndata.read_h5ad(data_filepath)

    # reach for raw transcript counts in the .raw attribute
    if hasattr(adata, "raw") and adata.raw is not None:
        adata.X = adata.raw.X
    print(f"Done loading data for {len(adata)} cells.")

    # re-index gene names
    adata.var["feature_name"] = adata.var["feature_name"].astype(str)
    adata.var["ensembl_ids"] = adata.var.index
    adata.var_names = adata.var["feature_name"]
    adata.var_names_make_unique(join="_")

    # filter cells & genes below minimum occurence threshold
    sc.pp.filter_cells(adata, min_genes=200)
    sc.pp.filter_genes(adata, min_cells=3)

    # annotate the group of mitochondrial genes as 'mt'
    adata.var["mt"] = adata.var_names.str.startswith("MT-")
    sc.pp.calculate_qc_metrics(
        adata, qc_vars=["mt"], percent_top=None, log1p=False, inplace=True
    )

    adata = adata[adata.obs.n_genes_by_counts < 2500, :]
    adata = adata[adata.obs.pct_counts_mt < 200, :]
    print(f"Done filtering cells, remaining data of shape {adata.shape}.")

    raw_X = np.copy(adata.X.toarray())
    norm_X, rank_norm_X = normalize_and_rank_transform(
        np.copy(adata.X.todense()), normalize=True
    )
    # update adata object with normalized expression
    adata.X = np.log10(1 + norm_X)

    # create dataframe of ranks and expression values for plotting
    expr_and_rank_df = pd.DataFrame(
        {
            "raw_transcript_count": np.ravel(raw_X),
            "preprocessed_transcript_count": np.ravel(norm_X),
            "preprocessed_rank": np.ravel(rank_norm_X),
            "log_preprocessed_transcript_count": np.log10(1 + np.ravel(norm_X)),
            "log_preprocessed_rank": np.log10(1 + np.ravel(rank_norm_X)),
        }
    )
    # remove 0 expression entries in the cellxgene matrix
    expr_and_rank_df = expr_and_rank_df[expr_and_rank_df["raw_transcript_count"] != 0]
    print(f"Done normalizing data, {len(expr_and_rank_df)} data points remaining.")

    # compute metrics for transformation to cells and back
    evaluate_transformation(df=expr_and_rank_df, plotting_sample_size=10000)

    preprocessed_output_filepath = data_filepath.parent / (
        data_filepath.stem + data_filepath.suffix.replace(".h5ad", "_preprocessed.h5ad")
    )
    print(f"Saving preprocessed transcript counts to {preprocessed_output_filepath}.")
    adata.write_h5ad(preprocessed_output_filepath)

    # convert the adata into ranked sequences of gene names ("cell sentences")
    csdata = utils.csdata_from_adata(adata)

    # make text files containing the cell sentences
    txt_output_dir = output_dir / "cell_sentences"
    txt_output_dir.mkdir(exist_ok=True, parents=True)
    utils.xlm_prepare_outpath(csdata, txt_output_dir, species_tag="human")
    print(f"Done writing cell sentences to file.")

    # make arrow-formatted dataset compatible with HuggingFace's datasets
    hf_output_dir = output_dir / "cell_sentences_hf"
    hf_output_dir.mkdir(exist_ok=True, parents=True)
    data_splits = ["train", "valid", "test"]
    data_files = {
        data_split: str(txt_output_dir / f"{data_split}_human.txt")
        for data_split in data_splits
    }
    dataset = load_dataset("text", data_files=data_files)

    # load cell type labels if available with transcript counts
    for data_split in data_splits:
        dataset[data_split] = dataset[data_split].rename_column("text", "input_ids")
        # retrieve split chunk from preprocessed transcript counts
        dataset_split_sample_indices = np.load(
            txt_output_dir / f"{data_split}_partition_indices.npy"
        )
        adata_split = adata[dataset_split_sample_indices, :].copy()
        if "cell_type" in adata_split.obs.columns:
            cell_type_labels = {"cell_type": adata_split.obs["cell_type"].tolist()}
            cell_type_dataset = Dataset.from_dict(cell_type_labels)
            dataset[data_split] = concatenate_datasets(
                [dataset[data_split], cell_type_dataset], axis=1
            )

    dataset.save_to_disk(hf_output_dir)
    print(f"Done transforming data to cell sentences.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data_filepath",
        type=Path,
        help="Input data filepath.",
        default=utils.DATA_DIR / "dominguez_sample.h5ad",
    )
    parser.add_argument(
        "--output_dir",
        type=Path,
        help="Output directory filepath.",
        default=utils.DATA_DIR,
    )
    args = parser.parse_args()

    main(args.data_filepath, args.output_dir)

transform.py
vandijklab/cell2sentence-ft