This model understands the "language of biology" by treating gene expression lists as sentences. It achieves 89% accuracy on distinguishing 44 distinct cell types in human colon tissue, outperforming the base model significantly on this specific task.
The model was evaluated on a held-out test set of 500 cells.
Use the following code to predict cell types. Note that the prompt format is critical for good performance.
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4# Load model
5model_name = "Jyx0208/C2S-UC-Gemma-2B"
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
8
9def predict_cell_type(cell_sentence, top_k=200):
10 """
11 Args:
12 cell_sentence: Space-separated string of gene names ordered by expression level (descending)
13 """
14 # Create prompt (Cell2Sentence Standard Format)
15 num_genes = len(cell_sentence.split()[:top_k])
16 prompt = f"The following is a list of {num_genes} gene names ordered by descending expression level in a human cell. Your task is to give the cell type which this cell belongs to based on its gene expression.\nCell sentence: {cell_sentence}.\nThe cell type corresponding to these genes is:"
17
18 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
19
20 with torch.no_grad():
21 outputs = model.generate(
22 **inputs,
23 max_new_tokens=50,
24 do_sample=False
25 )
26
27 answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
28 # Extract prediction (remove prompt)
29 prediction = answer[len(prompt):].strip().split('\n')[0]
30 return prediction
31
32# Example: Plasma IgA Cell
33# (Top expressed genes: JCHAIN, IGHA1, IGHA2, MZB1, ...)
34example_cell = "JCHAIN IGHA1 IGHA2 MZB1 TNFRSF17 SSR4 TXNDC5 FKBP11 SEC11C"
35print(f"Prediction: {predict_cell_type(example_cell)}")
1@article{cell2sentence2024,
2 title={Cell2Sentence: Teaching Large Language Models the Language of Biology},
3 author={Levine, Daniel and Rizvi, Syed and others},
4 journal={bioRxiv},
5 year={2024}
6}
This model is fine-tuned from Gemma and is subject to the
Gemma Terms of Use.