149 lines
6.4 KiB
Markdown
149 lines
6.4 KiB
Markdown
---
|
|
license: cc-by-4.0
|
|
datasets:
|
|
- vandijklab/immune-c2s
|
|
language:
|
|
- en
|
|
tags:
|
|
- pytorch
|
|
- causal-lm
|
|
- scRNA-seq
|
|
---
|
|
# Overview
|
|
|
|
This is the the Pythia-160m model developed by EleutherAI fine-tuned using Cell2Sentence on *full* scRNA-seq cells.
|
|
Cell2Sentence is a novel method for adapting large language models to single-cell transcriptomics.
|
|
We transform single-cell RNA sequencing data into sequences of gene names ordered by expression level, termed "cell sentences".
|
|
For more details, we refer to the paper linked below.
|
|
This model was trained on the immune tissue dataset from [Domínguez et al.](https://www.science.org/doi/10.1126/science.abl5197)
|
|
using 8 A100 40GB GPUs for approximately 20 hours on the following tasks:
|
|
1. conditional cell generation
|
|
2. unconditional cell generation
|
|
3. cell type prediction
|
|
|
|
## Cell2Sentence Links:
|
|
GitHub: <https://github.com/vandijklab/cell2sentence-ft> (Note: Codebase has Apache 2.0 license, weights shared on HuggingFace are CC-by-4.0)
|
|
Paper: <https://www.biorxiv.org/content/10.1101/2023.09.11.557287v3>
|
|
|
|
## Pythia Links:
|
|
GitHub: <https://github.com/EleutherAI/pythia>
|
|
Paper: <https://arxiv.org/abs/2304.01373>
|
|
Hugging Face: <https://huggingface.co/EleutherAI/pythia-160m>
|
|
|
|
# Evaluation
|
|
|
|
This model was evaluated on KNN classification and Gromov-Wasserstein (GW) distance.
|
|
The label for a generated cell is the corresponding cell type used in its corresponding prompt for generation.
|
|
Ground truth cells were sampled with replacement from a held out test dataset.
|
|
The generated cells are converted to expression vectors using the method described in the paper.
|
|
For complete details on the experiments, we refer to the paper.
|
|
|
|
| Model | k=3 NN (↑) | k=5 NN (↑) | k=10 NN (↑) | k=25 NN (↑) | GW (↓) |
|
|
| :---- | :---: | :---: | :---: | :---: | :----: |
|
|
| scGEN | 0.2376 | 0.2330 | 0.2377 | 0.2335 | 315.9505 |
|
|
| scVI | 0.2436 | 0.2400 | 0.2425 | 0.2348 | 302.1285 |
|
|
| scDiffusion | 0.2335 | 0.2288 | 0.2368 | 0.2306 | 72.0208 |
|
|
| scGPT | 0.1838 | 0.1788 | 0.1811 | 0.1882 | 2989.8066 |
|
|
| **C2S (Pythia-160m)** | **0.2588** | **0.2565** | **0.2746** | **0.2715** | **54.3040** |
|
|
|
|
# Sample Code
|
|
|
|
We provide an example of how to use the model to conditionally generate a cell equipped with a post-processing function to remove duplicate and invalid genes.
|
|
In order to generate full cells, the `max_length` generation parameter should be changed to 9200.
|
|
However, we recommend using an A100 GPU for inference speed and memory capacity if full cell generation is required.
|
|
Unconditional cell generation and cell type prediction prompts are included as well, but we do not include an example cell sentence to format the prompt.
|
|
We refer to the paper and GitHub repository for instructions on how to transform expression vectors into cell sentences.
|
|
|
|
```
|
|
import json
|
|
import re
|
|
from collections import Counter
|
|
from typing import List
|
|
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
|
|
|
|
def post_process_generated_cell_sentences(
|
|
cell_sentence: str,
|
|
gene_dictionary: List
|
|
):
|
|
"""
|
|
Post-processing function for generated cell sentences.
|
|
Invalid genes are removed and ranks of duplicated genes are averaged.
|
|
|
|
Arguments:
|
|
cell_sentence: generated cell sentence string
|
|
gene_dictionary: list of gene vocabulary (all uppercase)
|
|
|
|
Returns:
|
|
post_processed_sentence: generated cell sentence after post processing steps
|
|
"""
|
|
generated_gene_names = cell_sentence.split(" ")
|
|
generated_gene_names = [generated_gene.upper() for generated_gene in generated_gene_names]
|
|
|
|
#--- Remove nonsense genes ---#
|
|
generated_gene_names = [gene_name for gene_name in generated_gene_names if gene_name in gene_dictionary]
|
|
|
|
#--- 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
|
|
|
|
genes_path = "pbmc_vocab.json"
|
|
|
|
with open(vocab_path, "r") as f:
|
|
gene_dictionary = json.load(f)
|
|
|
|
model_name = "vandijklab/pythia-160m-c2s"
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_name,
|
|
torch_dtype=torch.float16,
|
|
attn_implementation="flash_attention_2"
|
|
).to(torch.device("cuda"))
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
|
|
cell_type = "T Cell"
|
|
ccg = f"Enumerate the genes in a {cell_type} cell with nonzero expression, from highest to lowest."
|
|
|
|
# Prompts for other forms a generation.
|
|
# ucg = "Display a cell's genes by expression level, in descending order."
|
|
# cellsentence = "CELL_SENTENCE"
|
|
# ctp = "Identify the cell type most likely associated with these highly expressed genes listed in descending order. "
|
|
# + cellsentence +
|
|
# "Name the cell type connected to these genes, ranked from highest to lowest expression."
|
|
|
|
tokens = tokenizer(ccg, return_tensors='pt')
|
|
input_ids = tokens['input_ids'].to(torch.device("cuda"))
|
|
attention_mask = tokens['attention_mask'].to(torch.device("cuda"))
|
|
|
|
with torch.no_grad():
|
|
outputs = model.generate(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
do_sample=True,
|
|
max_length=1024,
|
|
top_k=50,
|
|
top_p=0.95,
|
|
)
|
|
|
|
output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
cell_sentence = "".join(re.split(r"\?|\.|:", output_text)[1:]).strip()
|
|
processed_genes = post_process_generated_cell_sentences(cell_sentence, gene_dictionary)
|
|
``` |