Эффективно извлекайте последовательности для скользящего окна (большой набор данных)

Набор данных, который я сохранил, представляет собой просто координаты последовательности ДНК.

дф:

chr   start    stop   label  
chr1   9000    9100    1  
chr1   8803    8903     1  
chr1   8903    9000     0  

Моя цель — расширить исходный набор данных, создав скользящее окно вокруг каждой координаты для захвата контекстных последовательностей.

новый_дф:

chr    start         stop        label  
chr1   9000-5000    9000+5000      1  
chr1   9001-5000    9001+5000      1  
chr1   9002-5000    9002+5000      1  
...
chr1   9100-5000    9100+5000      1  
...

1. Создать скользящее окно. Для каждой записи в моем исходном наборе данных сгенерируйте новые строки, создав скользящее окно из 5000 нуклеотидов по обе стороны от координат начала и остановки. Это эффективно фиксирует контекст последовательности вокруг каждой исходной координаты.
2. Развернуть набор данных: для каждой исходной записи создайте новые строки, где каждая строка представляет определенную позицию в контекстном окне из 5000 нуклеотидов. Например, если исходное начало — 9000, я сгенерирую строки со стартовыми значениями от 9000-5000 до 9000+5000 и аналогичным образом настрою конечные значения. Длина каждой последовательности теперь составляет 10 001 символ.

используя эту функцию:

def expand_coordinates(element_locs, context=3):
    # Vectorized expansion of coordinates
    start = element_locs['Start'].astype(int)
    end = element_locs['End'].astype(int)
    
    expanded_data = []
    for idx, row in element_locs.iterrows():
        chr_name = row['Chromosome']
        chr_start = start[idx]
        chr_end = end[idx]
        
        for i in range(chr_start, chr_end + 1):
            expanded_data.append({
                'Chromosome': chr_name,
                'Start': max((i - 1) - context, 0),
                'End': min(i + context, max_sizes[chr_name])
            })
    
    expanded_df = pd.DataFrame(expanded_data)
    return expanded_df

3. Извлечение последовательностей: используйте расширенные координаты из new_df, чтобы получить соответствующие последовательности ДНК из набора данных.

def get_element_seqs(element_locs, context=3):
    expanded_df = expand_coordinates(element_locs, context=context)
    # Optimize genome fetching
    genome = pysam.Fastafile(ref_genome)
    def fetch_sequences(row):
        return genome.fetch(row['Chromosome'], row['Start'], row['End'])
    # Fetch sequences in a vectorized way
    expanded_df['sequence'] = expanded_df.apply(fetch_sequences, axis=1)
    return element_seqs

4. Токенизировать последовательности

dataset = Dataset.from_pandas(element_final[['Chromosome', 'sequence', 'label']]) 

dataset = dataset.shuffle(seed=42)
tokenizer = AutoTokenizer.from_pretrained(f"InstaDeepAI/nucleotide-transformer-500m-human-ref")
def tokenize_function(examples):
    outputs = tokenizer.batch_encode_plus(examples["sequence"], return_tensors = "pt", truncation=False, padding=False, max_length=80)
    return outputs
    
# Creating tokenized  dataset
tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True, batch_size=2000)

5. Получить вложения из токенов

input_file = f"tokenized_elements/tokenized_{ELEMENT_LABEL}/{filename}.arrow"

# Load input data
d1 = Dataset.from_file(input_file)

def embed_function(examples):
    torch.cuda.empty_cache()
    gc.collect()

    inputs = torch.tensor(examples['input_ids'])  # Convert to tensor
    inputs = inputs.to(device)

    with torch.no_grad():
        outputs = model(input_ids=inputs, output_hidden_states=True)

    # Step 3: Extract the embeddings
    hidden_states = outputs.hidden_states  # List of hidden states from all layers
    embeddings = hidden_states[-1]  # Assuming you want embeddings from the last layer
    averaged_embeddings = torch.mean(embeddings, dim=1)  # Calculate mean along dimension 1 (the dimension with size 86)
    averaged_embeddings = averaged_embeddings.to(torch.float32)  # Ensure float32 data type
    return {'embeddings': averaged_embeddings}

# Map embeddings function to input data
embeddings = d1.map(embed_function, batched=True, batch_size=1550)
embeddings = embeddings.remove_columns(["input_ids", "attention_mask"])

# Save embeddings to disk
output_dir = f"embedded_elements/embeddings_{ELEMENT_LABEL}/{filename}"  # Assuming ELEMENT_LABEL is defined elsewhere

6. Подключите вложения к XGBoost.

В результате я получаю огромные наборы данных, что приводит к сбою моего кода (например, я начинаю с 700 тысяч строк, а они расширяются до 1000 миллионов строк). Я использую панд, так что, может быть, в этом тоже проблема? Другая проблема в том, что я, кажется, не использую пакетную обработку? К сожалению, мой код продолжает давать сбой между шагами 2+3. Я думаю, что мне нужно реализовать пакетную обработку, но я не уверен, как все будет работать, поскольку в конечном итоге мне нужно будет передать выходные данные в LLM.