Pythia 160m C2s

🚀 Pythia-160m-c2s模型

這是由EleutherAI開發的Pythia-160m模型，使用Cell2Sentence方法在完整的單細胞RNA測序（scRNA-seq）細胞上進行了微調。Cell2Sentence是一種將大語言模型應用於單細胞轉錄組學的新方法。我們將單細胞RNA測序數據轉換為由表達水平排序的基因名稱序列，稱為“細胞句子”。更多詳細信息，請參考下面鏈接的論文。該模型在來自Domínguez等人的免疫組織數據集上進行訓練，使用8塊A100 40GB GPU，耗時約20小時，完成以下任務：

1. 條件細胞生成
2. 無條件細胞生成
3. 細胞類型預測

🚀 快速開始

本模型是基於Pythia-160m模型，利用Cell2Sentence方法在完整的scRNA-seq細胞上微調得到。Cell2Sentence能將單細胞轉錄組數據轉化為“細胞句子”，便於大語言模型處理。

✨ 主要特性

• 創新方法應用：採用Cell2Sentence方法，將單細胞RNA測序數據轉換為基因名稱序列，適配大語言模型處理單細胞轉錄組學。
• 多任務訓練：在條件細胞生成、無條件細胞生成和細胞類型預測等任務上進行訓練。
• 高性能表現：在KNN分類和Gromov - Wasserstein（GW）距離評估中表現出色。

📦 安裝指南

文檔未提供安裝步驟，此處跳過。

💻 使用示例

基礎用法

我們提供了一個如何使用該模型進行條件細胞生成的示例，其中包含一個後處理函數，用於去除重複和無效的基因。為了生成完整的細胞，需要將max_length生成參數更改為9200。不過，如果需要生成完整細胞，建議使用A100 GPU以保證推理速度和內存容量。

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)

高級用法

文檔未提及高級用法相關內容，此處跳過。

📚 詳細文檔

評估指標

本模型在KNN分類和Gromov - Wasserstein（GW）距離上進行了評估。生成細胞的標籤是其生成提示中對應的細胞類型。真實細胞是從保留的測試數據集中有放回採樣得到的。生成的細胞使用論文中描述的方法轉換為表達向量。完整的實驗細節請參考論文。

相關鏈接

Cell2Sentence鏈接：

• GitHub: https://github.com/vandijklab/cell2sentence-ft
• 論文: https://www.biorxiv.org/content/10.1101/2023.09.11.557287v3

Pythia鏈接：

• GitHub: https://github.com/EleutherAI/pythia
• 論文: https://arxiv.org/abs/2304.01373
• Hugging Face: https://huggingface.co/EleutherAI/pythia-160m

🔧 技術細節

文檔未提供足夠的技術實現細節，此處跳過。

📄 許可證

本項目採用CC BY - NC - ND 4.0許可證。