:zap: Ready the app.py for next experiment. Looks ok

app.py

@@ -1,21 +1,17 @@
+import logging
 import os
-import random
-
-import huggingface_hub
-import numpy as np
-from datasets import load_dataset, Dataset
-from dotenv import load_dotenv
-from pytorch_lightning import LightningDataModule
-from pytorch_lightning.utilities.types import TRAIN_DATALOADERS, EVAL_DATALOADERS
-from torch.utils.data import DataLoader, IterableDataset
-from sklearn.metrics import accuracy_score, recall_score, precision_score, f1_score, roc_auc_score
-# from torchmetrics.classification import BinaryAccuracy, BinaryAUROC, BinaryF1Score, BinaryPrecision, BinaryRecall
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModel, BertModel
-from transformers import TrainingArguments, Trainer
+from typing import Any
+
+from huggingface_hub import PyTorchModelHubMixin
+from pytorch_lightning import Trainer, LightningModule, LightningDataModule
+from pytorch_lightning.utilities.types import OptimizerLRScheduler, STEP_OUTPUT, EVAL_DATALOADERS
+from torch.utils.data import DataLoader, Dataset, IterableDataset
+from torcheval.metrics import BinaryAccuracy, BinaryAUROC, BinaryF1Score, BinaryPrecision, BinaryRecall
+from transformers import BertModel, BatchEncoding, BertTokenizer, TrainingArguments
+from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions
 import torch
-import logging
-import wandb
-
+from torch import nn
+from datasets import load_dataset
 timber = logging.getLogger()
 # logging.basicConfig(level=logging.DEBUG)
 logging.basicConfig(level=logging.INFO)  # change to level=logging.DEBUG to print more logs...
@@ -32,36 +28,273 @@ white = "\u001b[37m"
 FORWARD = "FORWARD_INPUT"
 BACKWARD = "BACKWARD_INPUT"
 
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+DNA_BERT_6 = "zhihan1996/DNA_bert_6"
+
 
-PRETRAINED_MODEL_NAME: str = "zhihan1996/DNA_bert_6"
+class CommonAttentionLayer(nn.Module):
+  def __init__(self, hidden_size, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self.attention_linear = nn.Linear(hidden_size, 1)
+    pass
+
+  def forward(self, hidden_states):
+    # Apply linear layer
+    attn_weights = self.attention_linear(hidden_states)
+    # Apply softmax to get attention scores
+    attn_weights = torch.softmax(attn_weights, dim=1)
+    # Apply attention weights to hidden states
+    context_vector = torch.sum(attn_weights * hidden_states, dim=1)
+    return context_vector, attn_weights
 
 
-def insert_debug_motif_at_random_position(seq, DEBUG_MOTIF):
-  start = 0
-  end = len(seq)
-  rand_pos = random.randrange(start, (end - len(DEBUG_MOTIF)))
-  random_end = rand_pos + len(DEBUG_MOTIF)
-  output = seq[start: rand_pos] + DEBUG_MOTIF + seq[random_end: end]
-  assert len(seq) == len(output)
-  return output
+class ReshapedBCEWithLogitsLoss(nn.BCEWithLogitsLoss):
+  def forward(self, input, target):
+    return super().forward(input.squeeze(), target.float())
 
 
-class PagingMQTLDataset(IterableDataset):
+class MQtlDnaBERT6Classifier(nn.Module, PyTorchModelHubMixin):
   def __init__(self,
-               m_dataset,
-               seq_len,
-               tokenizer,
-               max_length=512,
-               check_if_pipeline_is_ok_by_inserting_debug_motif=False):
-    self.dataset = m_dataset
-    self.check_if_pipeline_is_ok_by_inserting_debug_motif = check_if_pipeline_is_ok_by_inserting_debug_motif
-    self.debug_motif = "ATCGCCTA"
-    self.seq_len = seq_len
+               bert_model=BertModel.from_pretrained(pretrained_model_name_or_path=DNA_BERT_6),
+               hidden_size=768,
+               num_classes=1,
+               *args,
+               **kwargs
+               ):
+    super().__init__(*args, **kwargs)
+
+    self.model_name = "MQtlDnaBERT6Classifier"
+
+    self.bert_model = bert_model
+    self.attention = CommonAttentionLayer(hidden_size)
+    self.classifier = nn.Linear(hidden_size, num_classes)
+    pass
+
+  def forward(self, input_ids: torch.tensor, attention_mask: torch.tensor, token_type_ids):
+    """
+    # torch.Size([128, 1, 512]) --> [128, 512]
+    input_ids = input_ids.squeeze(dim=1).to(DEVICE)
+    # torch.Size([16, 1, 512]) --> [16, 512]
+    attention_mask = attention_mask.squeeze(dim=1).to(DEVICE)
+    token_type_ids = token_type_ids.squeeze(dim=1).to(DEVICE)
+    """
+    bert_output: BaseModelOutputWithPoolingAndCrossAttentions = self.bert_model(
+      input_ids=input_ids,
+      attention_mask=attention_mask,
+      token_type_ids=token_type_ids
+    )
+
+    last_hidden_state = bert_output.last_hidden_state
+    context_vector, ignore_attention_weight = self.attention(last_hidden_state)
+    y = self.classifier(context_vector)
+    return y
+
+"""
+class TorchMetrics:
+  def __init__(self):
+    self.binary_accuracy = BinaryAccuracy()  #.to(device)
+    self.binary_auc = BinaryAUROC()  # .to(device)
+    self.binary_f1_score = BinaryF1Score()  # .to(device)
+    self.binary_precision = BinaryPrecision()  # .to(device)
+    self.binary_recall = BinaryRecall()  # .to(device)
+    pass
+
+  def update_on_each_step(self, batch_predicted_labels, batch_actual_labels):  # todo: Add log if needed
+    # it looks like the library maintainers changed preds to input, ie, before: preds, now: input
+    self.binary_accuracy.update(input=batch_predicted_labels, target=batch_actual_labels)
+    self.binary_auc.update(input=batch_predicted_labels, target=batch_actual_labels)
+    self.binary_f1_score.update(input=batch_predicted_labels, target=batch_actual_labels)
+    self.binary_precision.update(input=batch_predicted_labels, target=batch_actual_labels)
+    self.binary_recall.update(input=batch_predicted_labels, target=batch_actual_labels)
+    pass
+
+  def compute_and_log_on_each_step(self, log, log_prefix: str, log_color: str = green):
+    b_accuracy = self.binary_accuracy.compute()
+    b_auc = self.binary_auc.compute()
+    b_f1_score = self.binary_f1_score.compute()
+    b_precision = self.binary_precision.compute()
+    b_recall = self.binary_recall.compute()
+    timber.info(log_color + f"{log_prefix}_acc = {b_accuracy}, {log_prefix}_auc = {b_auc}, {log_prefix}_f1_score = {b_f1_score}, {log_prefix}_precision = {b_precision}, {log_prefix}_recall = {b_recall}")
+    log(f"{log_prefix}_accuracy", b_accuracy)
+    log(f"{log_prefix}_auc", b_auc)
+    log(f"{log_prefix}_f1_score", b_f1_score)
+    log(f"{log_prefix}_precision", b_precision)
+    log(f"{log_prefix}_recall", b_recall)
+
+  # def reset_on_epoch_end(self):
+  #   self.binary_accuracy.reset()
+  #   self.binary_auc.reset()
+  #   self.binary_f1_score.reset()
+  #   self.binary_precision.reset()
+  #   self.binary_recall.reset()
+
+  def compute_and_reset_on_epoch_end(self, log, log_prefix: str, log_color: str = green):
+    b_accuracy = self.binary_accuracy.compute()
+    b_auc = self.binary_auc.compute()
+    b_f1_score = self.binary_f1_score.compute()
+    b_precision = self.binary_precision.compute()
+    b_recall = self.binary_recall.compute()
+    timber.info(  log_color + f"{log_prefix}_acc = {b_accuracy}, {log_prefix}_auc = {b_auc}, {log_prefix}_f1_score = {b_f1_score}, {log_prefix}_precision = {b_precision}, {log_prefix}_recall = {b_recall}")
+    log(f"{log_prefix}_accuracy", b_accuracy)
+    log(f"{log_prefix}_auc", b_auc)
+    log(f"{log_prefix}_f1_score", b_f1_score)
+    log(f"{log_prefix}_precision", b_precision)
+    log(f"{log_prefix}_recall", b_recall)
+
+    self.binary_accuracy.reset()
+    self.binary_auc.reset()
+    self.binary_f1_score.reset()
+    self.binary_precision.reset()
+    self.binary_recall.reset()
+    pass
+"""
+
+
+class TorchMetrics:
+  def __init__(self):
+    self.binary_accuracy = BinaryAccuracy()  #.to(device)
+    self.binary_auc = BinaryAUROC()  # .to(device)
+    self.binary_f1_score = BinaryF1Score()  # .to(device)
+    self.binary_precision = BinaryPrecision()  # .to(device)
+    self.binary_recall = BinaryRecall()  # .to(device)
+    pass
+
+  def update_on_each_step(self, batch_predicted_labels, batch_actual_labels):  # todo: Add log if needed
+    # it looks like the library maintainers changed preds to input, ie, before: preds, now: input
+    self.binary_accuracy.update(input=batch_predicted_labels, target=batch_actual_labels)
+    self.binary_auc.update(input=batch_predicted_labels, target=batch_actual_labels)
+    self.binary_f1_score.update(input=batch_predicted_labels, target=batch_actual_labels)
+    self.binary_precision.update(input=batch_predicted_labels, target=batch_actual_labels)
+    self.binary_recall.update(input=batch_predicted_labels, target=batch_actual_labels)
+    pass
+
+  def compute_and_reset_on_epoch_end(self, log, log_prefix: str, log_color: str = green):
+    b_accuracy = self.binary_accuracy.compute()
+    b_auc = self.binary_auc.compute()
+    b_f1_score = self.binary_f1_score.compute()
+    b_precision = self.binary_precision.compute()
+    b_recall = self.binary_recall.compute()
+    timber.info(  log_color + f"{log_prefix}_acc = {b_accuracy}, {log_prefix}_auc = {b_auc}, {log_prefix}_f1_score = {b_f1_score}, {log_prefix}_precision = {b_precision}, {log_prefix}_recall = {b_recall}")
+    log(f"{log_prefix}_accuracy", b_accuracy)
+    log(f"{log_prefix}_auc", b_auc)
+    log(f"{log_prefix}_f1_score", b_f1_score)
+    log(f"{log_prefix}_precision", b_precision)
+    log(f"{log_prefix}_recall", b_recall)
+
+    self.binary_accuracy.reset()
+    self.binary_auc.reset()
+    self.binary_f1_score.reset()
+    self.binary_precision.reset()
+    self.binary_recall.reset()
+    pass
+
+
+
+class MQtlBertClassifierLightningModule(LightningModule):
+  def __init__(self,
+               classifier: nn.Module,
+               criterion=None,  # nn.BCEWithLogitsLoss(),
+               regularization: int = 2,  # 1 == L1, 2 == L2, 3 (== 1 | 2) == both l1 and l2, else ignore / don't care
+               l1_lambda=0.001,
+               l2_wright_decay=0.001,
+               *args: Any,
+               **kwargs: Any):
+    super().__init__(*args, **kwargs)
+    self.classifier = classifier
+    self.criterion = criterion
+    self.train_metrics = TorchMetrics()
+    self.validate_metrics = TorchMetrics()
+    self.test_metrics = TorchMetrics()
+
+    self.regularization = regularization
+    self.l1_lambda = l1_lambda
+    self.l2_weight_decay = l2_wright_decay
+    pass
+
+  def forward(self, x, *args: Any, **kwargs: Any) -> Any:
+    input_ids: torch.tensor = x["input_ids"]
+    attention_mask: torch.tensor = x["attention_mask"]
+    token_type_ids: torch.tensor = x["token_type_ids"]
+    # print(f"\n{ type(input_ids) = }, {input_ids = }")
+    # print(f"{ type(attention_mask) = }, { attention_mask = }")
+    # print(f"{ type(token_type_ids) = }, { token_type_ids = }")
+
+    return self.classifier.forward(input_ids, attention_mask, token_type_ids)
+
+  def configure_optimizers(self) -> OptimizerLRScheduler:
+    # Here we add weight decay (L2 regularization) to the optimizer
+    weight_decay = 0.0
+    if self.regularization == 2 or self.regularization == 3:
+      weight_decay = self.l2_weight_decay
+    return torch.optim.Adam(self.parameters(), lr=1e-3, weight_decay=weight_decay)  # , weight_decay=0.005)
+
+  def training_step(self, batch, batch_idx, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+    # Accuracy on training batch data
+    x, y = batch
+    preds = self.forward(x)
+    loss = self.criterion(preds, y)
+
+    if self.regularization == 1 or self.regularization == 3:  # apply l1 regularization
+      l1_norm = sum(p.abs().sum() for p in self.parameters())
+      loss += self.l1_lambda * l1_norm
+
+    self.log("train_loss", loss)
+    # calculate the scores start
+    self.train_metrics.update_on_each_step(batch_predicted_labels=preds.squeeze(), batch_actual_labels=y)
+    # self.train_metrics.compute_and_log_on_each_step(log=self.log, log_prefix="train")
+    # calculate the scores end
+    return loss
+
+  def on_train_epoch_end(self) -> None:
+    self.train_metrics.compute_and_reset_on_epoch_end(log=self.log, log_prefix="train")
+    pass
 
+  def validation_step(self, batch, batch_idx, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+    # Accuracy on validation batch data
+    # print(f"debug { batch = }")
+    x, y = batch
+    preds = self.forward(x)
+    loss = self.criterion(preds, y)
+    """ loss = 0 #  <------------------------- maybe the loss calculation is problematic """
+    self.log("valid_loss", loss)
+    # calculate the scores start
+    self.validate_metrics.update_on_each_step(batch_predicted_labels=preds.squeeze(), batch_actual_labels=y)
+    # self.validate_metrics.compute_and_log_on_each_step(log=self.log, log_prefix="validate", log_color=blue)
+
+    # calculate the scores end
+    return loss
+
+  def on_validation_epoch_end(self) -> None:
+    self.validate_metrics.compute_and_reset_on_epoch_end(log=self.log, log_prefix="validate", log_color=blue)
+    return None
+
+  def test_step(self, batch, batch_idx, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+    # Accuracy on validation batch data
+    x, y = batch
+    preds = self.forward(x)
+    loss = self.criterion(preds, y)
+    self.log("test_loss", loss)  # do we need this?
+    # calculate the scores start
+    self.test_metrics.update_on_each_step(batch_predicted_labels=preds.squeeze(), batch_actual_labels=y)
+    # self.test_metrics.compute_and_log_on_each_step(log=self.log, log_prefix="test", log_color=magenta)
+
+    # calculate the scores end
+    return loss
+
+  def on_test_epoch_end(self) -> None:
+    self.test_metrics.compute_and_reset_on_epoch_end(log=self.log, log_prefix="test", log_color=magenta)
+    return None
+
+  pass
+
+
+class PagingMQTLDnaBertDataset(IterableDataset):
+  def __init__(self, dataset, tokenizer, max_length=512):  # hold on! why is it 512? I added 4000, and it crashed, the error suggested 512, that's why 512
+    self.dataset = dataset
     self.bert_tokenizer = tokenizer
     self.max_length = max_length
-    pass
+
+  # def __len__(self):
+  #   return len(self.dataset)
 
   def __iter__(self):
     for row in self.dataset:
@@ -71,267 +304,131 @@ class PagingMQTLDataset(IterableDataset):
 
   def preprocess(self, row):
     sequence = row['sequence']  # Fetch the 'sequence' column
-    if len(sequence) != self.seq_len:
-      return None  # skip problematic row!
     label = row['label']  # Fetch the 'label' column (or whatever target you use)
-    if label == 1 and self.check_if_pipeline_is_ok_by_inserting_debug_motif:
-      sequence = insert_debug_motif_at_random_position(seq=sequence, DEBUG_MOTIF=self.debug_motif)
 
-    input_ids = self.bert_tokenizer(sequence)["input_ids"]
-    tokenized_tensor = torch.tensor(input_ids)
-    label_tensor = torch.tensor(label)
-    output_dict = {"input_ids": tokenized_tensor, "labels": label_tensor}  # so this is now you do it?
-    return output_dict  # tokenized_tensor, label_tensor
+    # Tokenize the sequence
+    encoded_sequence: BatchEncoding = self.bert_tokenizer(
+      sequence,
+      truncation=True,
+      padding='max_length',
+      max_length=self.max_length,
+      return_tensors='pt'
+    )
+
+    encoded_sequence_squeezed = {key: val.squeeze() for key, val in encoded_sequence.items()}
+    return encoded_sequence_squeezed, label
 
 
-class MqtlDataModule(LightningDataModule):
-  def __init__(self, train_ds, val_ds, test_ds, batch_size=16):
+class DNABERTDataModule(LightningDataModule):
+  def __init__(self, model_name=DNA_BERT_6, batch_size=8, WINDOW=-1, is_local=False):
     super().__init__()
+    self.tokenized_dataset = None
+    self.dataset = None
+    self.train_dataset: PagingMQTLDnaBertDataset = None
+    self.validate_dataset: PagingMQTLDnaBertDataset = None
+    self.test_dataset: PagingMQTLDnaBertDataset = None
+    self.tokenizer = BertTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)
     self.batch_size = batch_size
-    self.train_loader = DataLoader(train_ds, batch_size=self.batch_size, shuffle=False,
-                                   # collate_fn=collate_fn,
-                                   num_workers=1,
-                                   # persistent_workers=True
-                                   )
-    self.validate_loader = DataLoader(val_ds, batch_size=self.batch_size, shuffle=False,
-                                      # collate_fn=collate_fn,
-                                      num_workers=1,
-                                      # persistent_workers=True
-                                      )
-    self.test_loader = DataLoader(test_ds, batch_size=self.batch_size, shuffle=False,
-                                  # collate_fn=collate_fn,
-                                  num_workers=1,
-                                  # persistent_workers=True
-                                  )
-    pass
+    self.is_local = is_local
+    self.window = WINDOW
 
   def prepare_data(self):
-    pass
-
-  def setup(self, stage: str) -> None:
-    timber.info(f"inside setup: {stage = }")
-    pass
+    # Download and prepare dataset
+    data_files = {
+      # small samples
+      "train_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_train_binned.csv",
+      "validate_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_validate_binned.csv",
+      "test_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_test_binned.csv",
+      # medium samples
+      "train_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_train_binned.csv",
+      "validate_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_validate_binned.csv",
+      "test_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_test_binned.csv",
+
+      # large samples
+      "train_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_train_binned.csv",
+      "validate_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_validate_binned.csv",
+      "test_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_test_binned.csv",
+
+      # really tiny
+      # "tiny_train": "/home/soumic/Codes/mqtl-classification/src/inputdata/tiny_dataset_4000_train_binned.csv",
+      # "tiny_validate": "/home/soumic/Codes/mqtl-classification/src/inputdata/tiny_dataset_4000_validate_binned.csv",
+      # "tiny_test": "/home/soumic/Codes/mqtl-classification/src/inputdata/tiny_dataset_4000_test_binned.csv",
+
+    }
+    if self.is_local:
+      self.dataset = load_dataset("csv", data_files=data_files, streaming=True)
+    else:
+      self.dataset = load_dataset("fahimfarhan/mqtl-classification-datasets")
+
+  def setup(self, stage=None):
+    self.train_dataset = PagingMQTLDnaBertDataset(self.dataset['train_binned_4000'], self.tokenizer)
+    self.validate_dataset = PagingMQTLDnaBertDataset(self.dataset['validate_binned_4000'], self.tokenizer)
+    self.test_dataset = PagingMQTLDnaBertDataset(self.dataset['test_binned_4000'], self.tokenizer)
+
+  def train_dataloader(self):
+    return DataLoader(self.train_dataset, batch_size=self.batch_size, num_workers=1)
+
+  def val_dataloader(self):
+    return DataLoader(self.validate_dataset, batch_size=self.batch_size, num_workers=1)
 
-  def train_dataloader(self) -> TRAIN_DATALOADERS:
-    return self.train_loader
+  def test_dataloader(self) -> EVAL_DATALOADERS:
+    return DataLoader(self.test_dataset, batch_size=self.batch_size, num_workers=1)
 
-  def val_dataloader(self) -> EVAL_DATALOADERS:
-    return self.validate_loader
 
-  def test_dataloader(self) -> EVAL_DATALOADERS:
-    return self.test_loader
-
-
-def create_paging_train_val_test_datasets(tokenizer, WINDOW, is_debug, batch_size=1000):
-  data_files = {
-    # small samples
-    "train_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_train_binned.csv",
-    "validate_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_validate_binned.csv",
-    "test_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_test_binned.csv",
-    # medium samples
-    "train_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_train_binned.csv",
-    "validate_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_validate_binned.csv",
-    "test_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_test_binned.csv",
-
-    # large samples
-    "train_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_train_binned.csv",
-    "validate_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_validate_binned.csv",
-    "test_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_test_binned.csv",
-  }
-
-  dataset_map = None
+def start_bert(classifier_model, model_save_path, criterion, WINDOW, batch_size=4,
+               dataset_folder_prefix="inputdata/", is_binned=True, is_debug=False, max_epochs=10):
   is_my_laptop = os.path.isfile("/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_train_binned.csv")
-  if is_my_laptop:
-    dataset_map = load_dataset("csv", data_files=data_files, streaming=True)
-  else:
-    dataset_map = load_dataset("fahimfarhan/mqtl-classification-datasets", streaming=True)
-
-  train_dataset = PagingMQTLDataset(dataset_map[f"train_binned_{WINDOW}"],
-                                    check_if_pipeline_is_ok_by_inserting_debug_motif=is_debug,
-                                    tokenizer=tokenizer,
-                                    seq_len=WINDOW
-                                    )
-  val_dataset = PagingMQTLDataset(dataset_map[f"validate_binned_{WINDOW}"],
-                                  check_if_pipeline_is_ok_by_inserting_debug_motif=is_debug,
-                                  tokenizer=tokenizer,
-                                  seq_len=WINDOW)
-  test_dataset = PagingMQTLDataset(dataset_map[f"test_binned_{WINDOW}"],
-                                   check_if_pipeline_is_ok_by_inserting_debug_motif=is_debug,
-                                   tokenizer=tokenizer,
-                                   seq_len=WINDOW)
-  # data_module = MqtlDataModule(train_ds=train_dataset, val_ds=val_dataset, test_ds=test_dataset, batch_size=batch_size)
-  return train_dataset, val_dataset, test_dataset
-
-
-def login_inside_huggingface_virtualmachine():
-  # Load the .env file, but don't crash if it's not found (e.g., in Hugging Face Space)
-  try:
-    load_dotenv()  # Only useful on your laptop if .env exists
-    print(".env file loaded successfully.")
-  except Exception as e:
-    print(f"Warning: Could not load .env file. Exception: {e}")
-
-  # Try to get the token from environment variables
-  try:
-    token = os.getenv("HF_TOKEN")
-
-    if not token:
-      raise ValueError("HF_TOKEN not found. Make sure to set it in the environment variables or .env file.")
-
-    # Log in to Hugging Face Hub
-    huggingface_hub.login(token)
-    print("Logged in to Hugging Face Hub successfully.")
-
-  except Exception as e:
-    print(f"Error during Hugging Face login: {e}")
-    # Handle the error appropriately (e.g., exit or retry)
-
-  # wand db login
-  try:
-    api_key = os.getenv("WAND_DB_API_KEY")
-    timber.info(f"{api_key = }")
-
-    if not api_key:
-      raise ValueError("WAND_DB_API_KEY not found. Make sure to set it in the environment variables or .env file.")
-
-    # Log in to Hugging Face Hub
-    wandb.login(key=api_key)
-    print("Logged in to wand db successfully.")
-
-  except Exception as e:
-    print(f"Error during wand db Face login: {e}")
-  pass
+  model_local_directory = f"my-awesome-model-{WINDOW}"
+  model_remote_repository = f"fahimfarhan/dnabert-6-mqtl-classifier-{WINDOW}"
+  file_suffix = ""
+  if is_binned:
+    file_suffix = "_binned"
 
+  data_module = DNABERTDataModule(batch_size=batch_size, WINDOW=WINDOW, is_local=is_my_laptop)
 
-# use sklearn cz torchmetrics.classification gave array index out of bound exception :/ (whatever it is called in python)
-def compute_metrics_using_sklearn(p):
-  try:
-    pred, labels = p
+  # classifier_model = classifier_model.to(DEVICE)
 
-    # Get predicted class labels
-    pred_labels = np.argmax(pred, axis=1)
+  classifier_module = MQtlBertClassifierLightningModule(
+    classifier=classifier_model,
+    regularization=2, criterion=criterion)
 
-    # Get predicted probabilities for the positive class
-    pred_probs = pred[:, 1]  # Assuming binary classification and 2 output classes
+  # if os.path.exists(model_save_path):
+  #   classifier_module.load_state_dict(torch.load(model_save_path))
 
-    accuracy = accuracy_score(y_true=labels, y_pred=pred_labels)
-    recall = recall_score(y_true=labels, y_pred=pred_labels)
-    precision = precision_score(y_true=labels, y_pred=pred_labels)
-    f1 = f1_score(y_true=labels, y_pred=pred_labels)
-    roc_auc = roc_auc_score(y_true=labels, y_score=pred_probs)
+  classifier_module = classifier_module  # .double()
 
-    return {"accuracy": accuracy, "roc_auc": roc_auc, "precision": precision, "recall": recall, "f1": f1}
+  # Prepare data using the DataModule
+  data_module.prepare_data()
+  data_module.setup()
 
-  except Exception as x:
-    print(f"compute_metrics_using_sklearn failed with exception: {x}")
-    return {"accuracy": 0, "roc_auc": 0, "precision": 0, "recall": 0, "f1": 0}
+  trainer = Trainer(max_epochs=max_epochs, precision="32")
 
+  # Train the model
+  trainer.fit(model=classifier_module, datamodule=data_module)
+  trainer.test(model=classifier_module, datamodule=data_module)
+  torch.save(classifier_module.state_dict(), model_save_path)
 
-def start():
-  os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+  # classifier_module.push_to_hub("fahimfarhan/mqtl-classifier-model")
 
-  login_inside_huggingface_virtualmachine()
-  WINDOW = 4000
-  batch_size = 100
-  model_local_directory = f"my-awesome-model-{WINDOW}"
-  model_remote_repository = f"fahimfarhan/dnabert-6-mqtl-classifier-{WINDOW}"
-
-  is_my_laptop = os.path.isfile("/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_train_binned.csv")
+  classifier_model.save_pretrained(save_directory=model_local_directory, safe_serialization=False)
+  # push to the hub
+  commit_message = f":tada: Push model for window size {WINDOW} from huggingface space"
+  if is_my_laptop:
+    commit_message = f":tada: Push model for window size {WINDOW} from zephyrus"
 
-  tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL_NAME, trust_remote_code=True)
-  classifier_model = AutoModelForSequenceClassification.from_pretrained(PRETRAINED_MODEL_NAME, num_labels=2)
-  args = {
-    "output_dir": "output_dnabert-6-mqtl_classification",
-    "num_train_epochs": 1,
-    "max_steps": 20_000,  # train 36k + val 4k = 40k
-    # Set the number of steps you expect to train, originally 1000, takes too much time. So I set it to 10 to run faster and check my code/pipeline
-    "run_name": "laptop_run_dna-bert-6-mqtl_classification",  # Override run_name here
-    "per_device_train_batch_size": 1,
-    "gradient_accumulation_steps": 32,
-    "gradient_checkpointing": True,
-    "learning_rate": 1e-3,
-    "save_safetensors": False,  # I added it. this solves the runtime error!
-    # not sure if it is a good idea. sklearn may slow down training, causing time loss... if so, disable these 2 lines below
-    "evaluation_strategy": "epoch",  # To calculate metrics per epoch
-    "logging_strategy": "epoch"  # Extra: to log training data stats for loss
-  }
-
-  training_args = TrainingArguments(**args)
-  # train_dataset, eval_dataset, test_dataset = create_data_module(tokenizer=tokenizer, WINDOW=WINDOW,
-  #                                                                batch_size=batch_size,
-  #                                                                is_debug=False)
-  """  # example code
-  max_length = 32_000
-  sequence = 'ACTG' * int(max_length / 4)
-  # sequence = 'ACTG' * int(1000) # seq_len = 4000 it works!
-  sequence = [sequence] * 8  # Create 8 identical samples
-  tokenized = tokenizer(sequence)["input_ids"]
-  labels = [0, 1] * 4
-
-  # Create a dataset for training
-  run_the_code_ds = Dataset.from_dict({"input_ids": tokenized, "labels": labels})
-  run_the_code_ds.set_format("pt")
-  """
-
-  train_ds, val_ds, test_ds = create_paging_train_val_test_datasets(tokenizer, WINDOW=WINDOW, is_debug=False)
-  # train_ds, val_ds, test_ds = run_the_code_ds, run_the_code_ds, run_the_code_ds
-  # train_ds.set_format("pt") # doesn't work!
-
-  trainer = Trainer(
-    model=classifier_model,
-    args=training_args,
-    train_dataset=train_ds,
-    eval_dataset=val_ds,
-    compute_metrics=compute_metrics_using_sklearn  # torch_metrics.compute_metrics
+  classifier_model.push_to_hub(
+    repo_id=model_remote_repository,
+    # subfolder=f"my-awesome-model-{WINDOW}", subfolder didn't work :/
+    commit_message=commit_message,  # f":tada: Push model for window size {WINDOW}"
+    # safe_serialization=False
   )
-  # train, and validate
-  result = trainer.train()
-  try:
-    print(f"{result = }")
-  except Exception as x:
-    print(f"{x = }")
-
-  # testing
-  try:
-    # with torch.no_grad(): # didn't work :/
-    test_results = trainer.evaluate(eval_dataset=test_ds)
-    print(f"{test_results = }")
-  except Exception as oome:
-    print(f"{oome = }")
-  finally:
-    # save the model
-    model_name = "DnaBert6MQtlClassifier"
-
-    classifier_model.save_pretrained(save_directory=model_local_directory, safe_serialization=False)
-
-    # push to the hub
-    commit_message = f":tada: Push model for window size {WINDOW} from huggingface space"
-    if is_my_laptop:
-      commit_message = f":tada: Push model for window size {WINDOW} from zephyrus"
-
-    classifier_model.push_to_hub(
-      repo_id=model_remote_repository,
-      # subfolder=f"my-awesome-model-{WINDOW}", subfolder didn't work :/
-      commit_message=commit_message,  # f":tada: Push model for window size {WINDOW}"
-      safe_serialization=False
-    )
   pass
 
 
-def interprete_demo():
-  is_my_laptop = True
-  WINDOW = 4000
-  batch_size = 100
-  model_local_directory = f"my-awesome-model-{WINDOW}"
-  model_remote_repository = f"fahimfarhan/dnabert-6-mqtl-classifier-{WINDOW}"
-
-  try:
-    classifier_model = AutoModel.from_pretrained(model_remote_repository)
-    # todo: use captum / gentech-grelu to interpret the model
-  except Exception as x:
-    print(x)
-
-
-if __name__ == '__main__':
-  start()
+if __name__ == "__main__":
+  dataset_folder_prefix = "inputdata/"
+  pytorch_model = MQtlDnaBERT6Classifier()
+  start_bert(classifier_model=pytorch_model, model_save_path=f"weights_{pytorch_model.model_name}.pth",
+             criterion=ReshapedBCEWithLogitsLoss(), WINDOW=4000, batch_size=12,  # max 14 on my laptop...
+             dataset_folder_prefix=dataset_folder_prefix, max_epochs=1)
   pass

app_v1_backup.py

@@ -1,5 +1,7 @@
+import os
 from typing import Any
 
+from huggingface_hub import PyTorchModelHubMixin
 from pytorch_lightning import Trainer, LightningModule, LightningDataModule
 from pytorch_lightning.utilities.types import OptimizerLRScheduler, STEP_OUTPUT, EVAL_DATALOADERS
 from torch.utils.data import DataLoader, Dataset
@@ -46,7 +48,7 @@ class ReshapedBCEWithLogitsLoss(nn.BCEWithLogitsLoss):
     return super().forward(input.squeeze(), target.float())
 
 
-class MQtlDnaBERT6Classifier(nn.Module):
+class MQtlDnaBERT6Classifier(nn.Module, PyTorchModelHubMixin):
   def __init__(self,
                bert_model=BertModel.from_pretrained(pretrained_model_name_or_path=DNA_BERT_6),
                hidden_size=768,
@@ -277,6 +279,10 @@ data_module = DNABERTDataModule(model_name=model_name, batch_size=8)
 
 def start_bert(classifier_model, model_save_path, criterion, WINDOW=200, batch_size=4,
                dataset_folder_prefix="inputdata/", is_binned=True, is_debug=False, max_epochs=10):
+
+  is_my_laptop = os.path.isfile("/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_train_binned.csv")
+  model_local_directory = f"my-awesome-model-{WINDOW}"
+  model_remote_repository = f"fahimfarhan/dnabert-6-mqtl-classifier-{WINDOW}"
   file_suffix = ""
   if is_binned:
     file_suffix = "_binned"
@@ -324,7 +330,20 @@ def start_bert(classifier_model, model_save_path, criterion, WINDOW=200, batch_s
   trainer.test(model=classifier_module, datamodule=data_module)
   torch.save(classifier_module.state_dict(), model_save_path)
 
-  classifier_module.push_to_hub("fahimfarhan/mqtl-classifier-model")
+  # classifier_module.push_to_hub("fahimfarhan/mqtl-classifier-model")
+
+  classifier_model.save_pretrained(save_directory=model_local_directory, safe_serialization=False)
+  # push to the hub
+  commit_message = f":tada: Push model for window size {WINDOW} from huggingface space"
+  if is_my_laptop:
+    commit_message = f":tada: Push model for window size {WINDOW} from zephyrus"
+
+  classifier_model.push_to_hub(
+    repo_id=model_remote_repository,
+    # subfolder=f"my-awesome-model-{WINDOW}", subfolder didn't work :/
+    commit_message=commit_message,  # f":tada: Push model for window size {WINDOW}"
+    safe_serialization=False
+  )
   pass
 
 
@@ -332,6 +351,6 @@ if __name__ == "__main__":
   dataset_folder_prefix = "inputdata/"
   pytorch_model = MQtlDnaBERT6Classifier()
   start_bert(classifier_model=pytorch_model, model_save_path=f"weights_{pytorch_model.model_name}.pth",
-             criterion=ReshapedBCEWithLogitsLoss(), WINDOW=200, batch_size=4,
+             criterion=ReshapedBCEWithLogitsLoss(), WINDOW=4000, batch_size=4,
              dataset_folder_prefix=dataset_folder_prefix, max_epochs=2)
   pass

failed_app.py

@@ -0,0 +1,337 @@
+import os
+import random
+
+import huggingface_hub
+import numpy as np
+from datasets import load_dataset, Dataset
+from dotenv import load_dotenv
+from pytorch_lightning import LightningDataModule
+from pytorch_lightning.utilities.types import TRAIN_DATALOADERS, EVAL_DATALOADERS
+from torch.utils.data import DataLoader, IterableDataset
+from sklearn.metrics import accuracy_score, recall_score, precision_score, f1_score, roc_auc_score
+# from torchmetrics.classification import BinaryAccuracy, BinaryAUROC, BinaryF1Score, BinaryPrecision, BinaryRecall
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModel, BertModel
+from transformers import TrainingArguments, Trainer
+import torch
+import logging
+import wandb
+
+timber = logging.getLogger()
+# logging.basicConfig(level=logging.DEBUG)
+logging.basicConfig(level=logging.INFO)  # change to level=logging.DEBUG to print more logs...
+
+black = "\u001b[30m"
+red = "\u001b[31m"
+green = "\u001b[32m"
+yellow = "\u001b[33m"
+blue = "\u001b[34m"
+magenta = "\u001b[35m"
+cyan = "\u001b[36m"
+white = "\u001b[37m"
+
+FORWARD = "FORWARD_INPUT"
+BACKWARD = "BACKWARD_INPUT"
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+PRETRAINED_MODEL_NAME: str = "zhihan1996/DNA_bert_6"
+
+
+def insert_debug_motif_at_random_position(seq, DEBUG_MOTIF):
+  start = 0
+  end = len(seq)
+  rand_pos = random.randrange(start, (end - len(DEBUG_MOTIF)))
+  random_end = rand_pos + len(DEBUG_MOTIF)
+  output = seq[start: rand_pos] + DEBUG_MOTIF + seq[random_end: end]
+  assert len(seq) == len(output)
+  return output
+
+
+class PagingMQTLDataset(IterableDataset):
+  def __init__(self,
+               m_dataset,
+               seq_len,
+               tokenizer,
+               max_length=512,
+               check_if_pipeline_is_ok_by_inserting_debug_motif=False):
+    self.dataset = m_dataset
+    self.check_if_pipeline_is_ok_by_inserting_debug_motif = check_if_pipeline_is_ok_by_inserting_debug_motif
+    self.debug_motif = "ATCGCCTA"
+    self.seq_len = seq_len
+
+    self.bert_tokenizer = tokenizer
+    self.max_length = max_length
+    pass
+
+  def __iter__(self):
+    for row in self.dataset:
+      processed = self.preprocess(row)
+      if processed is not None:
+        yield processed
+
+  def preprocess(self, row):
+    sequence = row['sequence']  # Fetch the 'sequence' column
+    if len(sequence) != self.seq_len:
+      return None  # skip problematic row!
+    label = row['label']  # Fetch the 'label' column (or whatever target you use)
+    if label == 1 and self.check_if_pipeline_is_ok_by_inserting_debug_motif:
+      sequence = insert_debug_motif_at_random_position(seq=sequence, DEBUG_MOTIF=self.debug_motif)
+
+    input_ids = self.bert_tokenizer(sequence)["input_ids"]
+    tokenized_tensor = torch.tensor(input_ids)
+    label_tensor = torch.tensor(label)
+    output_dict = {"input_ids": tokenized_tensor, "labels": label_tensor}  # so this is now you do it?
+    return output_dict  # tokenized_tensor, label_tensor
+
+
+class MqtlDataModule(LightningDataModule):
+  def __init__(self, train_ds, val_ds, test_ds, batch_size=16):
+    super().__init__()
+    self.batch_size = batch_size
+    self.train_loader = DataLoader(train_ds, batch_size=self.batch_size, shuffle=False,
+                                   # collate_fn=collate_fn,
+                                   num_workers=1,
+                                   # persistent_workers=True
+                                   )
+    self.validate_loader = DataLoader(val_ds, batch_size=self.batch_size, shuffle=False,
+                                      # collate_fn=collate_fn,
+                                      num_workers=1,
+                                      # persistent_workers=True
+                                      )
+    self.test_loader = DataLoader(test_ds, batch_size=self.batch_size, shuffle=False,
+                                  # collate_fn=collate_fn,
+                                  num_workers=1,
+                                  # persistent_workers=True
+                                  )
+    pass
+
+  def prepare_data(self):
+    pass
+
+  def setup(self, stage: str) -> None:
+    timber.info(f"inside setup: {stage = }")
+    pass
+
+  def train_dataloader(self) -> TRAIN_DATALOADERS:
+    return self.train_loader
+
+  def val_dataloader(self) -> EVAL_DATALOADERS:
+    return self.validate_loader
+
+  def test_dataloader(self) -> EVAL_DATALOADERS:
+    return self.test_loader
+
+
+def create_paging_train_val_test_datasets(tokenizer, WINDOW, is_debug, batch_size=1000):
+  data_files = {
+    # small samples
+    "train_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_train_binned.csv",
+    "validate_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_validate_binned.csv",
+    "test_binned_200": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_200_test_binned.csv",
+    # medium samples
+    "train_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_train_binned.csv",
+    "validate_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_validate_binned.csv",
+    "test_binned_1000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_1000_test_binned.csv",
+
+    # large samples
+    "train_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_train_binned.csv",
+    "validate_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_validate_binned.csv",
+    "test_binned_4000": "/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_test_binned.csv",
+  }
+
+  dataset_map = None
+  is_my_laptop = os.path.isfile("/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_train_binned.csv")
+  if is_my_laptop:
+    dataset_map = load_dataset("csv", data_files=data_files, streaming=True)
+  else:
+    dataset_map = load_dataset("fahimfarhan/mqtl-classification-datasets", streaming=True)
+
+  train_dataset = PagingMQTLDataset(dataset_map[f"train_binned_{WINDOW}"],
+                                    check_if_pipeline_is_ok_by_inserting_debug_motif=is_debug,
+                                    tokenizer=tokenizer,
+                                    seq_len=WINDOW
+                                    )
+  val_dataset = PagingMQTLDataset(dataset_map[f"validate_binned_{WINDOW}"],
+                                  check_if_pipeline_is_ok_by_inserting_debug_motif=is_debug,
+                                  tokenizer=tokenizer,
+                                  seq_len=WINDOW)
+  test_dataset = PagingMQTLDataset(dataset_map[f"test_binned_{WINDOW}"],
+                                   check_if_pipeline_is_ok_by_inserting_debug_motif=is_debug,
+                                   tokenizer=tokenizer,
+                                   seq_len=WINDOW)
+  # data_module = MqtlDataModule(train_ds=train_dataset, val_ds=val_dataset, test_ds=test_dataset, batch_size=batch_size)
+  return train_dataset, val_dataset, test_dataset
+
+
+def login_inside_huggingface_virtualmachine():
+  # Load the .env file, but don't crash if it's not found (e.g., in Hugging Face Space)
+  try:
+    load_dotenv()  # Only useful on your laptop if .env exists
+    print(".env file loaded successfully.")
+  except Exception as e:
+    print(f"Warning: Could not load .env file. Exception: {e}")
+
+  # Try to get the token from environment variables
+  try:
+    token = os.getenv("HF_TOKEN")
+
+    if not token:
+      raise ValueError("HF_TOKEN not found. Make sure to set it in the environment variables or .env file.")
+
+    # Log in to Hugging Face Hub
+    huggingface_hub.login(token)
+    print("Logged in to Hugging Face Hub successfully.")
+
+  except Exception as e:
+    print(f"Error during Hugging Face login: {e}")
+    # Handle the error appropriately (e.g., exit or retry)
+
+  # wand db login
+  try:
+    api_key = os.getenv("WAND_DB_API_KEY")
+    timber.info(f"{api_key = }")
+
+    if not api_key:
+      raise ValueError("WAND_DB_API_KEY not found. Make sure to set it in the environment variables or .env file.")
+
+    # Log in to Hugging Face Hub
+    wandb.login(key=api_key)
+    print("Logged in to wand db successfully.")
+
+  except Exception as e:
+    print(f"Error during wand db Face login: {e}")
+  pass
+
+
+# use sklearn cz torchmetrics.classification gave array index out of bound exception :/ (whatever it is called in python)
+def compute_metrics_using_sklearn(p):
+  try:
+    pred, labels = p
+
+    # Get predicted class labels
+    pred_labels = np.argmax(pred, axis=1)
+
+    # Get predicted probabilities for the positive class
+    pred_probs = pred[:, 1]  # Assuming binary classification and 2 output classes
+
+    accuracy = accuracy_score(y_true=labels, y_pred=pred_labels)
+    recall = recall_score(y_true=labels, y_pred=pred_labels)
+    precision = precision_score(y_true=labels, y_pred=pred_labels)
+    f1 = f1_score(y_true=labels, y_pred=pred_labels)
+    roc_auc = roc_auc_score(y_true=labels, y_score=pred_probs)
+
+    return {"accuracy": accuracy, "roc_auc": roc_auc, "precision": precision, "recall": recall, "f1": f1}
+
+  except Exception as x:
+    print(f"compute_metrics_using_sklearn failed with exception: {x}")
+    return {"accuracy": 0, "roc_auc": 0, "precision": 0, "recall": 0, "f1": 0}
+
+
+def start():
+  os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+  login_inside_huggingface_virtualmachine()
+  WINDOW = 4000
+  batch_size = 100
+  model_local_directory = f"my-awesome-model-{WINDOW}"
+  model_remote_repository = f"fahimfarhan/dnabert-6-mqtl-classifier-{WINDOW}"
+
+  is_my_laptop = os.path.isfile("/home/soumic/Codes/mqtl-classification/src/inputdata/dataset_4000_train_binned.csv")
+
+  tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL_NAME, trust_remote_code=True)
+  classifier_model = AutoModelForSequenceClassification.from_pretrained(PRETRAINED_MODEL_NAME, num_labels=2)
+  args = {
+    "output_dir": "output_dnabert-6-mqtl_classification",
+    "num_train_epochs": 1,
+    "max_steps": 20_000,  # train 36k + val 4k = 40k
+    # Set the number of steps you expect to train, originally 1000, takes too much time. So I set it to 10 to run faster and check my code/pipeline
+    "run_name": "laptop_run_dna-bert-6-mqtl_classification",  # Override run_name here
+    "per_device_train_batch_size": 1,
+    "gradient_accumulation_steps": 32,
+    "gradient_checkpointing": True,
+    "learning_rate": 1e-3,
+    "save_safetensors": False,  # I added it. this solves the runtime error!
+    # not sure if it is a good idea. sklearn may slow down training, causing time loss... if so, disable these 2 lines below
+    "evaluation_strategy": "epoch",  # To calculate metrics per epoch
+    "logging_strategy": "epoch"  # Extra: to log training data stats for loss
+  }
+
+  training_args = TrainingArguments(**args)
+  # train_dataset, eval_dataset, test_dataset = create_data_module(tokenizer=tokenizer, WINDOW=WINDOW,
+  #                                                                batch_size=batch_size,
+  #                                                                is_debug=False)
+  """  # example code
+  max_length = 32_000
+  sequence = 'ACTG' * int(max_length / 4)
+  # sequence = 'ACTG' * int(1000) # seq_len = 4000 it works!
+  sequence = [sequence] * 8  # Create 8 identical samples
+  tokenized = tokenizer(sequence)["input_ids"]
+  labels = [0, 1] * 4
+
+  # Create a dataset for training
+  run_the_code_ds = Dataset.from_dict({"input_ids": tokenized, "labels": labels})
+  run_the_code_ds.set_format("pt")
+  """
+
+  train_ds, val_ds, test_ds = create_paging_train_val_test_datasets(tokenizer, WINDOW=WINDOW, is_debug=False)
+  # train_ds, val_ds, test_ds = run_the_code_ds, run_the_code_ds, run_the_code_ds
+  # train_ds.set_format("pt") # doesn't work!
+
+  trainer = Trainer(
+    model=classifier_model,
+    args=training_args,
+    train_dataset=train_ds,
+    eval_dataset=val_ds,
+    compute_metrics=compute_metrics_using_sklearn  # torch_metrics.compute_metrics
+  )
+  # train, and validate
+  result = trainer.train()
+  try:
+    print(f"{result = }")
+  except Exception as x:
+    print(f"{x = }")
+
+  # testing
+  try:
+    # with torch.no_grad(): # didn't work :/
+    test_results = trainer.evaluate(eval_dataset=test_ds)
+    print(f"{test_results = }")
+  except Exception as oome:
+    print(f"{oome = }")
+  finally:
+    # save the model
+    model_name = "DnaBert6MQtlClassifier"
+
+    classifier_model.save_pretrained(save_directory=model_local_directory, safe_serialization=False)
+
+    # push to the hub
+    commit_message = f":tada: Push model for window size {WINDOW} from huggingface space"
+    if is_my_laptop:
+      commit_message = f":tada: Push model for window size {WINDOW} from zephyrus"
+
+    classifier_model.push_to_hub(
+      repo_id=model_remote_repository,
+      # subfolder=f"my-awesome-model-{WINDOW}", subfolder didn't work :/
+      commit_message=commit_message,  # f":tada: Push model for window size {WINDOW}"
+      safe_serialization=False
+    )
+  pass
+
+
+def interprete_demo():
+  is_my_laptop = True
+  WINDOW = 4000
+  batch_size = 100
+  model_local_directory = f"my-awesome-model-{WINDOW}"
+  model_remote_repository = f"fahimfarhan/dnabert-6-mqtl-classifier-{WINDOW}"
+
+  try:
+    classifier_model = AutoModel.from_pretrained(model_remote_repository)
+    # todo: use captum / gentech-grelu to interpret the model
+  except Exception as x:
+    print(x)
+
+
+if __name__ == '__main__':
+  start()
+  pass