Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+*.ipynb
+data/
+logs/
+notebooks/
+output/
+run.sh
+wandb/

README.md

@@ -1,7 +1,105 @@
+# 🧚 MethFormer: A Transformer for DNA Methylation
+
+**MethFormer** is a masked regression transformer model trained to learn local and long-range patterns in DNA methylation (5mC and 5hmC) across genomic regions. Pretrained on binned methylation data, it is designed for downstream fine-tuning on tasks such as predicting MLL binding or chromatin state.
+
 ---
-license: mit
-language:
-- en
-pipeline_tag: image-feature-extraction
-library_name: transformers
----
+
+## 🚀 Overview
+
+* **Inputs**: Binned methylation values (5mC, 5hmC) over 1024bp windows (32 bins × 2 channels)
+* **Pretraining objective**: Masked methylation imputation (per-bin regression)
+* **Architecture**: Transformer encoder with linear projection head
+* **Downstream tasks**: MLL binding prediction, chromatin state inference, or enhancer classification
+
+---
+
+## 📁 Project Structure
+
+```
+.
+├── config/                       # config
+├── data/                         # Binned methylation datasets (HuggingFace format)
+├── output/                       # Pretrained models, logs, and checkpoints
+├── scripts/                      
+│   ├── methformer.py             # Model classes, data collator, 
+│   ├── pretrain_methformer.py    # Main training script
+│   └── finetune_mll.py           # (optional) downstream fine-tuning
+├── requirements.txt
+└── README.md
+```
+
+---
+
+## 👩‍💻 Pretraining MethFormer
+
+### Step 1: Prepare Dataset
+
+Preprocess 5mC and 5hmC data into 1024bp windows, binned into 32 bins × 2 features. Save using Hugging Face's `datasets.DatasetDict` format:
+
+```
+DatasetDict({
+  train: Dataset({
+    features: ['input_values', 'attention_mask', 'labels']
+  }),
+  validation: Dataset(...)
+})
+```
+
+### Step 2: Run Pretraining
+
+```bash
+python scripts/pretrain_methformer.py
+```
+
+Options can be customized inside the script or modified for sweep tuning. This will:
+
+* Train the model using masked regression loss
+* Evaluate on a held-out chromosome (e.g., `chr8`)
+* Log metrics to [Weights & Biases](https://wandb.ai)
+* Save the best model checkpoint
+
+---
+
+## 📊 Metrics
+
+* `masked_mse`: Mean squared error over unmasked positions
+* `masked_mae`: Mean absolute error
+
+---
+
+## 🧪 Fine-tuning on MLL Binding
+
+After pretraining:
+
+1. Replace the regression head with a scalar head for MLL prediction.
+2. Use a `Trainer` to fine-tune on log1p-transformed MLL-N RPKM values mean over 1kb regions.
+
+See `scripts/finetune_mll.py` for an example.
+
+---
+
+## 🔍 Visualizations & Interpretability
+
+You can run [Captum](https://captum.ai) or SHAP for:
+
+* Per-bin attribution of 5mC/5hmC to MLL binding
+* Visualizing what MethFormer attends to during fine-tuning
+
+---
+
+## 🛠️ Dependencies
+
+Key packages:
+
+* `transformers`
+* `datasets`
+* `wandb`
+* `torch`
+* `anndata`
+* `scikit-learn`
+
+---
+
+## 🧠 Acknowledgements
+
+* Built with inspiration from DNABERT, Grover, and vision transformers

config/pretrain_sweep_config.json

@@ -0,0 +1,17 @@
+{
+  "name": "methformer_pretrain_sweep",
+  "method": "bayes",
+  "metric": {"name": "eval/masked_mse", "goal": "minimize"},
+  "early_terminate": {
+    "type": "hyperband",
+    "min_iter": 4,
+    "eta": 2
+  },
+  "parameters": {
+    "masking_ratio": {"values": [0.1, 0.15, 0.2]},
+    "hidden_dim": {"values": [64, 128, 256]},
+    "num_hidden_layers": {"values": [6, 8, 12]},
+    "num_attention_heads": {"values": [4, 8]},
+    "hidden_dropout_prob": {"values": [0.1, 0.2, 0.3]}
+  }
+}

requirements.txt

@@ -0,0 +1,5 @@
+datasets  
+scikit-learn
+torch
+transformers 
+wandb

scripts/methformer.py

@@ -0,0 +1,126 @@
+import random
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+
+class MethformerDataset(Dataset):
+    """
+    Dataset that returns masked inputs, original labels, and attention masks.
+    """
+
+    def __init__(
+        self, data_tensor, chunk_size=128, mask_value=-1.0, masking_ratio=0.15
+    ):
+        self.data = data_tensor
+        self.n_samples, self.n_regions, self.n_channels = self.data.shape
+        self.chunk_size = min(chunk_size, self.n_regions)
+        self.mask_value = mask_value
+        self.masking_ratio = masking_ratio
+
+    def __len__(self):
+        return self.n_samples * (self.n_regions // self.chunk_size)
+
+    def __getitem__(self, idx):
+        sample_idx = idx % self.n_samples
+        chunk_start = random.randint(0, self.n_regions - self.chunk_size)
+        chunk = self.data[sample_idx, chunk_start : chunk_start + self.chunk_size, :]
+
+        x = torch.tensor(chunk, dtype=torch.float32)
+        mask = torch.rand(self.chunk_size) < self.masking_ratio
+        x_masked = x.clone()
+        x_masked[mask] = self.mask_value
+
+        return {"inputs": x_masked, "labels": x, "attention_mask": ~mask}
+
+
+class MethformerCollator:
+    def __init__(self, masking_ratio=0.15):
+        self.masking_ratio = masking_ratio
+
+    def __call__(self, batch):
+        def ensure_tensor(x):
+            if isinstance(x, torch.Tensor):
+                return x
+            return torch.tensor(x, dtype=torch.float32)
+
+        inputs = [ensure_tensor(item["inputs"]) for item in batch]
+        labels = [ensure_tensor(item["labels"]) for item in batch]
+        attention_mask = [
+            torch.tensor(item["attention_mask"], dtype=torch.bool) for item in batch
+        ]
+
+        inputs_tensor = torch.stack(inputs)
+        labels_tensor = torch.stack(labels)
+        attention_mask_tensor = torch.stack(attention_mask)
+
+        return {
+            "input_values": inputs_tensor,
+            "labels": labels_tensor,
+            "attention_mask": attention_mask_tensor,
+        }
+
+
+class Methformer(PreTrainedModel):
+    """
+    Masked Transformer model for methylation data.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.input_dim = getattr(config, "input_dim", 2)
+        hidden_dim = getattr(config, "hidden_dim", 128)
+        num_layers = config.num_hidden_layers
+        num_heads = config.num_attention_heads
+        dropout = config.hidden_dropout_prob
+        max_len = getattr(config, "max_position_embeddings", 1024)
+
+        self.embed = nn.Linear(self.input_dim, hidden_dim)
+        self.pos_embed = nn.Parameter(torch.randn(1, max_len, hidden_dim))
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_dim, nhead=num_heads, dropout=dropout, batch_first=True
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.output_head = nn.Linear(hidden_dim, self.input_dim)
+
+    def forward(self, input_values, attention_mask, labels=None):
+        x = self.embed(input_values)
+        x = x + self.pos_embed[:, : x.size(1), :].to(x.device)
+
+        attn_mask = ~attention_mask.bool()
+        x = self.encoder(x, src_key_padding_mask=attn_mask)
+        output = self.output_head(x)
+
+        loss = None
+        if labels is not None:
+            mask = attention_mask.unsqueeze(-1).expand_as(labels)
+            loss_fn = nn.MSELoss()
+            loss = loss_fn(output[mask], labels[mask])
+
+        return ModelOutput(loss=loss, last_hidden_state=output)
+
+
+class MethformerRegressor(PreTrainedModel):
+    """
+    Regression model that uses Methformer as the encoder.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.encoder = Methformer(config)
+        self.regression_head = nn.Linear(config.hidden_dim, 1)
+
+    def forward(self, input_values, attention_mask, labels=None):
+        x = self.encoder(input_values, attention_mask)
+        pooled = (x * attention_mask.unsqueeze(-1)).sum(1) / attention_mask.sum(
+            1, keepdim=True
+        )
+        logits = self.regression_head(pooled)
+        loss = None
+        if labels is not None:
+            loss = F.mse_loss(logits, labels)
+        return {"loss": loss, "logits": logits}

scripts/pretrain_methformer.py

@@ -0,0 +1,128 @@
+import datetime
+import os
+
+import torch
+import wandb
+from datasets import load_from_disk
+from sklearn.metrics import mean_absolute_error, mean_squared_error
+from transformers import (
+    EarlyStoppingCallback,
+    PretrainedConfig,
+    Trainer,
+    TrainingArguments,
+)
+
+from methformer import (
+    Methformer,
+    MethformerCollator,
+)
+
+run_name = f"mf_{datetime.datetime.now().strftime('%Y-%m-%d_%H%M')}"
+print(f"Run name: {run_name}")
+
+out_dir = "/home/ubuntu/project/MethFormer/output/methformer_pretrained/"
+os.makedirs(out_dir, exist_ok=True)
+
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
+dataset = load_from_disk("/home/ubuntu/project/MethFormer/data/methformer_pretrain_binned")
+train_dataset = dataset["train"].shuffle(seed=42)
+eval_dataset = dataset["validation"]
+
+data_collator = MethformerCollator()
+
+config = PretrainedConfig(
+    input_dim=2,
+    hidden_dim=128,
+    num_hidden_layers=12,
+    num_attention_heads=8,
+    hidden_dropout_prob=0.1,
+)
+
+model = Methformer(config)
+model.to(device)
+
+training_args = TrainingArguments(
+    run_name=run_name,
+    output_dir=os.path.join(out_dir, "checkpoints"),
+    eval_on_start=True,
+    per_device_train_batch_size=128,
+    per_device_eval_batch_size=256,
+    gradient_accumulation_steps=1,
+    max_grad_norm=1.0,
+    learning_rate=1e-5,
+    warmup_ratio=0.05,
+    lr_scheduler_type="cosine",
+    num_train_epochs=20,
+    logging_dir=os.path.join(out_dir, "logs"),
+    save_strategy="steps",
+    save_total_limit=1,
+    eval_strategy="steps",
+    logging_steps=1000,
+    eval_steps=1000,
+    save_steps=5000,
+    metric_for_best_model="masked_mse",
+    greater_is_better=False,
+    report_to="wandb",
+    disable_tqdm=False,
+    dataloader_num_workers=8,
+    remove_unused_columns=False,
+    fp16=not torch.backends.mps.is_available(),
+    load_best_model_at_end=True,
+    seed=42,
+)
+
+
+def compute_metrics(eval_preds):
+    logits, labels = eval_preds
+    logits = torch.tensor(logits)
+    labels = torch.tensor(labels)
+    mask = labels != -1.0
+    masked_logits = logits[mask].cpu.numpy()
+    masked_labels = labels[mask].cpu.numpy()
+    mse = mean_squared_error(masked_labels, masked_logits)
+    mae = mean_absolute_error(masked_labels, masked_logits)
+    return {
+        "masked_mse": mse,
+        "masked_mae": mae,
+    }
+
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    compute_metrics=compute_metrics,
+    data_collator=data_collator,
+    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
+)
+
+print("Starting training...")
+
+wandb.init(
+    group="methformer_pretrain",
+    job_type="pretrain_full",
+    name=run_name,
+    dir=out_dir,
+    reinit="finish_previous",
+    config=config.to_dict(),
+)
+
+trainer.train()
+print("Training complete. Saving model...")
+
+save_path = f"{out_dir}/model"
+os.makedirs(save_path, exist_ok=True)
+trainer.save_model(save_path)
+model.config.save_pretrained(save_path)
+print(f"Model saved to {save_path}")
+
+wandb.finish()

scripts/pretrain_sweep.py

@@ -0,0 +1,149 @@
+import datetime
+import json
+import os
+
+import torch
+import wandb
+from datasets import load_from_disk
+from transformers import (
+    EarlyStoppingCallback,
+    PretrainedConfig,
+    Trainer,
+    TrainingArguments,
+)
+
+from methformer import Methformer, MethformerCollator
+
+
+def compute_metrics(eval_preds):
+    logits, labels = eval_preds
+    logits = torch.tensor(logits)
+    labels = torch.tensor(labels)
+
+    # Only evaluate masked positions (label == -1.0 was masked during input)
+    mask = labels != -1.0
+
+    masked_mse = torch.mean((logits[mask] - labels[mask]) ** 2).item()
+    masked_mae = torch.mean(torch.abs(logits[mask] - labels[mask])).item()
+
+    return {
+        "masked_mse": masked_mse,
+        "masked_mae": masked_mae,
+    }
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
+dataset = load_from_disk("/home/ubuntu/project/MethFormer/data/methformer_pretrain_binned")
+train_dataset = dataset["train"].shuffle(seed=42)
+eval_dataset = dataset["validation"]
+
+
+def train():
+    wandb.init(
+        group="methformer_pretrain_sweep",
+        job_type="pretrain_sweep",
+        name=f"mf_{datetime.datetime.now().strftime('%Y-%m-%d_%H%M')}",
+        dir="/home/ubuntu/project/MethFormer/output/methformer_pretrain_sweep",
+        reinit="finish_previous",
+    )
+    config = wandb.config
+
+    run_name = f"mf_{datetime.datetime.now().strftime('%Y-%m-%d_%H%M')}"
+    out_dir = f"/home/ubuntu/project/MethFormer/output/methformer_pretrain_sweep/{run_name}"
+    os.makedirs(out_dir, exist_ok=True)
+
+    model_config = PretrainedConfig(
+        input_dim=2,
+        hidden_dim=config.hidden_dim,
+        num_hidden_layers=config.num_hidden_layers,
+        num_attention_heads=config.num_attention_heads,
+        hidden_dropout_prob=config.hidden_dropout_prob,
+    )
+
+    model = Methformer(model_config)
+    model.to(device)
+
+    training_args = TrainingArguments(
+        run_name=run_name,
+        output_dir=os.path.join(out_dir, "checkpoints"),
+        eval_on_start=True,
+        per_device_train_batch_size=128,
+        per_device_eval_batch_size=256,
+        gradient_accumulation_steps=1,
+        max_grad_norm=1.0,
+        learning_rate=1e-5,
+        warmup_ratio=0.05,
+        lr_scheduler_type="cosine",
+        num_train_epochs=20,
+        logging_dir=os.path.join(out_dir, "logs"),
+        save_strategy="steps",
+        save_total_limit=1,
+        eval_strategy="steps",
+        logging_steps=500,
+        eval_steps=5000,
+        save_steps=5000,
+        metric_for_best_model="masked_mse",
+        greater_is_better=False,
+        report_to="wandb",
+        disable_tqdm=False,
+        dataloader_num_workers=8,
+        remove_unused_columns=False,
+        fp16=not torch.backends.mps.is_available(),
+        load_best_model_at_end=True,
+        seed=42,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        compute_metrics=compute_metrics,
+        data_collator=MethformerCollator(masking_ratio=config.masking_ratio),
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
+    )
+
+    trainer.train()
+
+    # Save the final model
+    model.save_pretrained(os.path.join(out_dir, "model"))
+    model.config.save_pretrained(os.path.join(out_dir, "model"))
+
+
+with open("/home/ubuntu/project/MethFormer/config/pretrain_sweep_config.json", "r") as f:
+    sweep_config = json.load(f)
+
+sweep_id = wandb.sweep(
+    sweep=sweep_config,
+    project="MethFormer",
+)
+
+wandb.agent(sweep_id, train, count=20)
+
+# After the sweep
+api = wandb.Api()
+
+sweep_path = f"{wandb.run.entity}/{wandb.run.project}/{sweep_id}"
+sweep = api.sweep(sweep_path)
+
+# Filter only finished runs with masked_r2
+runs = [
+    run for run in sweep.runs if run.state == "finished" and "masked_r2" in run.summary
+]
+
+# Find best run by highest masked_r2
+best_run = max(runs, key=lambda r: r.summary["masked_r2"])
+
+# Save best config
+best_config = {k: v for k, v in best_run.config.items() if not k.startswith("_")}
+with open("/home/ubuntu/project/MethFormer/config/best_config.json", "w") as f:
+    json.dump(best_config, f, indent=2)
+
+print(f"Best run ID: {best_run.id}")
+print(f"Best masked_r2: {best_run.summary['masked_r2']}")