<jupyter_start><jupyter_text>We will leverage [Orion, a knowledge discovery tool for scientific publications](https://github.com/orion-search/orion), to get a subset of publications from [biorXiv](https://www.biorxiv.org/). Then, we will keep those related to COVID-19 as well as a random sample of non-COVID-19 publications using [MAG's Fields of Study](https://arxiv.org/abs/1805.12216).<jupyter_code># Read abstracts from Microsoft Academic Graph
df = pd.read_csv("drive/MyDrive/Colab Notebooks/mag_papers.csv")
df.head()
print(f"Number of Covid-19 publications: {df[df.is_Covid19==1].shape[0]}")
print(f"Number of non-Covid-19 publications: {df[df.is_Covid19!=1].shape[0]}")
# Set a seed with Pytorch Lightning
pl.seed_everything(42)

# Split the dataframe to training and evaluation sets
df_train, df_val = train_test_split(df, test_size=.05)

# Split the dataframe to training and test sets
df_train, df_test = train_test_split(df_train, test_size=.05)
class CovidDataset(torch.utils.data.Dataset):
  def __init__(self, data, tokenizer, max_token_len=128):
    self.data = data
    self.tokenizer = tokenizer
    self.max_token_len = max_token_len

  def __len__(self):
    return len(self.data)

  def __getitem__(self, index):
    # Grab a single row from the dataframe
    data_row = self.data.iloc[index]

    # Grab text and label
    abstract_text = data_row.abstract
    labels = torch.tensor(data_row.is_Covid19)

    encoding = tokenizer.encode_plus(
        abstract_text,
        add_special_tokens=True,
        max_length=self.max_token_len,
        return_token_type_ids=False,
        padding="max_length",
        truncation=True,
        return_attention_mask=True,
        return_tensors="pt"
    )

    return dict(
        abstract_text=abstract_text,
        input_ids=encoding["input_ids"].flatten(),
        attention_mask=encoding["attention_mask"].flatten(),
        labels=labels
    )
class CovidDataModule(pl.LightningDataModule):
  def __init__(self, df_train, df_val, df_test, tokenizer, batch_size=8, max_token_len=512):
    super().__init__()
    self.df_train = df_train
    self.df_val = df_val
    self.df_test = df_test
    self.tokenizer = tokenizer
    self.batch_size = batch_size
    self.max_token_len = max_token_len

  def setup(self):
    self.train_dataset = CovidDataset(self.df_train, self.tokenizer, self.max_token_len)
    self.val_dataset = CovidDataset(self.df_val, self.tokenizer, self.max_token_len)
    self.test_dataset = CovidDataset(self.df_val, self.tokenizer, self.max_token_len)

  def train_dataloader(self):
    return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=4)

  def val_dataloader(self):
    return DataLoader(self.val_dataset, batch_size=self.batch_size)

  def test_dataloader(self):
    return DataLoader(self.test_dataset, batch_size=self.batch_size)
# Hyperparameters
N_EPOCHS = 4
BATCH_SIZE = 16

# Pick a transformer from Huggingface
TRANSFORMER_MODEL_NAME = "bert-base-uncased"

# Instantiate the tokenizer
tokenizer = BertTokenizer.from_pretrained(TRANSFORMER_MODEL_NAME)

# Create a Pytorch Lightning DataModule
data_module = CovidDataModule(df_train, df_val, df_test, tokenizer, batch_size=BATCH_SIZE, max_token_len=128)
data_module.setup()
class CovidClassifier(pl.LightningModule):
  def __init__(self, n_classes, lr=2e-5, steps_per_epoch=None, n_epochs=None):
    super().__init__()

    self.bert_model = BertModel.from_pretrained(TRANSFORMER_MODEL_NAME, return_dict=True)
    self.n_classes = n_classes
    self.classifier = nn.Linear(self.bert_model.config.hidden_size,  self.n_classes)
    self.steps_per_epoch = steps_per_epoch
    self.n_epochs = n_epochs
    self.lr = lr

    self.criterion = nn.CrossEntropyLoss()
    # self.criterion = nn.BCELoss()

    self.save_hyperparameters()

  def forward(self, input_ids, attention_mask, labels=None):
    output = self.bert_model(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output)

    # output = torch.sigmoid(output)

    loss = 0
    if labels is not None:
      loss = self.criterion(output.view(-1, self.n_classes), labels.view(-1))
      # loss = self.criterion(output, labels)
    return loss, output

  def training_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("train_loss", loss, prog_bar=True)
    return {"loss":loss, "predictions":outputs, "labels":labels}

  def validation_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("val_loss", loss, prog_bar=True)
    return loss
  
  def test_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("test_loss", loss, prog_bar=True)
    return loss

  def configure_optimizers(self):
    optimizer = AdamW(self.parameters(), lr=self.lr)
    warmup_steps = self.steps_per_epoch // 3
    total_steps = self.steps_per_epoch * self.n_epochs - warmup_steps
    
    scheduler = get_linear_schedule_with_warmup(
        optimizer, warmup_steps, total_steps
    )
    return [optimizer], [scheduler]
# wandb.login()
model = CovidClassifier(n_classes=2, steps_per_epoch=len(df_train) // BATCH_SIZE, n_epochs=N_EPOCHS)

wandb_logger = WandbLogger(
    project="testing",
    save_code=False,
    tags=["covid-classifier"],
    reinit=True,
    )

wandb_logger.watch(model, log='all')
wandb.run.name
trainer = pl.Trainer(max_epochs=N_EPOCHS, gpus=1, progress_bar_refresh_rate=30, logger=wandb_logger)
trainer.fit(model, data_module)
trainer.test()
wandb.finish()
trainer.save_checkpoint("covid_classifier.ckpt")
new_model = CovidClassifier.load_from_checkpoint(n_classes=2, checkpoint_path="covid_classifier.ckpt")
new_model.freeze()
x = trainer.test()
x
(len(df_train) // BATCH_SIZE) * 1 * BATCH_SIZE
df_train.shape<jupyter_output><empty_output><jupyter_text>code scraps<jupyter_code>import torch
from torch import nn
import pytorch_lightning as pl
from torch.utils.data import Dataset, DataLoader
from transformers import BertModel, BertTokenizer, AdamW, get_linear_schedule_with_warmup
from sklearn.model_selection import train_test_split
import numpy as np
import pandas as pd
# Read abstracts from Microsoft Academic Graph
df = pd.read_csv("drive/MyDrive/Colab Notebooks/mag_papers.csv")
df.head()
# Split the dataframe to training and evaluation sets
pl.seed_everything(42)
df_train, df_val = train_test_split(df, test_size=.05)
print(f"% of Covid-19 publications: {df[df.is_Covid19==1].shape[0]}")
print(f"% of non-Covid-19 publications: {df[df.is_Covid19!=1].shape[0]}")
df_train.shape
# Shuffle the training data
df_train = df_train.sample(frac=1.)
sample_row = df.iloc[15]
sample_abstract = sample_row.abstract
sample_label = sample_row.is_Covid19

print(sample_abstract)
print()
print(sample_label)
# Instantiate the tokenizer
TRANSFORMER_MODEL_NAME = "bert-base-uncased"
tokenizer = BertTokenizer.from_pretrained(TRANSFORMER_MODEL_NAME)
encoding = tokenizer.encode_plus(
    sample_abstract, 
    add_special_tokens=True,
    max_length=512, 
    return_token_type_ids=True, 
    padding="max_length",
    return_attention_mask=True,
    return_tensors="pt"
    )
encoding.keys()
len(encoding["token_type_ids"].squeeze(0))
encoding["input_ids"].shape
encoding["attention_mask"]
type(encoding["input_ids"])
#Instantiate the tokenizer
TRANSFORMER_MODEL_NAME = "bert-base-uncased"
tokenizer = BertTokenizer.from_pretrained(TRANSFORMER_MODEL_NAME)

class CovidDataset(torch.utils.data.Dataset):
  def __init__(self, data, tokenizer, max_token_len=128):
    self.data = data
    self.tokenizer = tokenizer
    self.max_token_len = max_token_len

  def __len__(self):
    return len(self.data)

  def __getitem__(self, index):
    # Grab a single row from the dataframe
    data_row = self.data.iloc[index]

    # Grab text and label
    abstract_text = data_row.abstract
    labels = torch.tensor(data_row.is_Covid19)

    # print(torch.FloatTensor(labels))

    encoding = tokenizer.encode_plus(
        abstract_text,
        add_special_tokens=True,
        max_length=self.max_token_len,
        return_token_type_ids=False,
        padding="max_length",
        truncation=True,
        return_attention_mask=True,
        return_tensors="pt"
    )

    return dict(
        abstract_text=abstract_text,
        input_ids=encoding["input_ids"].flatten(),
        attention_mask=encoding["attention_mask"].flatten(),
        labels=labels
    )
train_dataset = CovidDataset(df_train, tokenizer, max_token_len=512)
sample_items = train_dataset[0]
sample_items.keys()
sample_items['labels']
sample_items['abstract_text']

# load model
bert_model = BertModel.from_pretrained(TRANSFORMER_MODEL_NAME, return_dict=True)
sample_items["input_ids"].unsqueeze(dim=0).shape
prediction = bert_model(sample_items["input_ids"].unsqueeze(dim=0), sample_items["attention_mask"].unsqueeze(dim=0))
prediction.last_hidden_state.shape
class CovidDataModule(pl.LightningDataModule):
  def __init__(self, df_train, df_val, tokenizer, batch_size=8, max_token_len=512):
    super().__init__()
    self.df_train = df_train
    self.df_val = df_val
    self.tokenizer = tokenizer
    self.batch_size = batch_size
    self.max_token_len = max_token_len

  def setup(self):
    self.train_dataset = CovidDataset(self.df_train, self.tokenizer, self.max_token_len)
    self.val_dataset = CovidDataset(self.df_val, self.tokenizer, self.max_token_len)
    self.test_dataset = CovidDataset(self.df_val, self.tokenizer, self.max_token_len)

  def train_dataloader(self):
    return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=4)

  def val_dataloader(self):
    return DataLoader(self.val_dataset, batch_size=self.batch_size)

  def test_dataloader(self):
    return DataLoader(self.val_dataset, batch_size=self.batch_size)

N_EPOCHS = 4
BATCH_SIZE = 12

data_module = CovidDataModule(df_train, df_val, tokenizer, batch_size=BATCH_SIZE, max_token_len=128)
data_module.setup()<jupyter_output><empty_output><jupyter_text>Modelling<jupyter_code>target = torch.ones([10, 64], dtype=torch.float32)  # 64 classes, batch size = 10
output = torch.full([10, 64], 1.5)  # A prediction (logit)
pos_weight = torch.ones([64])  # All weights are equal to 1
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
criterion(output, target)  # -log(sigmoid(1.5))
nn.BCEWithLogitsLoss
class CovidClassifier(pl.LightningModule):
  def __init__(self, n_classes, steps_per_epoch=None, n_epochs=None):
    super().__init__()

    self.bert_model = BertModel.from_pretrained(TRANSFORMER_MODEL_NAME, return_dict=True)
    self.n_classes = n_classes
    self.classifier = nn.Linear(self.bert_model.config.hidden_size,  self.n_classes)
    self.steps_per_epoch = steps_per_epoch
    self.n_epochs = n_epochs

    self.criterion = nn.CrossEntropyLoss()
    # self.criterion = nn.BCELoss()

  def forward(self, input_ids, attention_mask, labels=None):
    output = self.bert_model(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output)

    # output = torch.sigmoid(output)

    loss = 0
    if labels is not None:
      loss = self.criterion(output.view(-1, self.n_classes), labels.view(-1))
      # loss = self.criterion(output, labels)
    return loss, output

  def training_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("train_loss", loss, prog_bar=True, logger=False)
    return {"loss":loss, "predictions":outputs, "labels":labels}

  def validation_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("val_loss", loss, prog_bar=True, logger=False)
    return loss
  
  def test_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("test_loss", loss, prog_bar=True, logger=False)
    return loss

  def configure_optimizers(self):
    optimizer = AdamW(self.parameters(), lr=2e-5)
    warmup_steps = self.steps_per_epoch // 3
    total_steps = self.steps_per_epoch * self.n_epochs - warmup_steps
    
    scheduler = get_linear_schedule_with_warmup(
        optimizer, warmup_steps, total_steps
    )
    return [optimizer], [scheduler]
model = CovidClassifier(n_classes=2, steps_per_epoch=len(df_train) // BATCH_SIZE, n_epochs=N_EPOCHS)
# # predictions with the untrained model
# _, predictions = model(sample_items["input_ids"].unsqueeze(dim=0), sample_items["attention_mask"].unsqueeze(dim=0))
# predictions
trainer = pl.Trainer(max_epochs=N_EPOCHS, gpus=1, progress_bar_refresh_rate=30)
trainer.fit(model, data_module)
trainer.test()
# # predictions with the untrained model
# _, predictions = model(sample_items["input_ids"].unsqueeze(dim=0), sample_items["attention_mask"].unsqueeze(dim=0))
# print(predictions)
# softmax = nn.Softmax(-1)
# print(softmax(predictions))
trainer.save_checkpoint("covid_classifier.ckpt")
new_model = CovidClassifier.load_from_checkpoint(n_classes=2, checkpoint_path="covid_classifier.ckpt")
new_model.freeze()




sample_abstract = "COVID-19 is a deadly coronavirus that broke out last year."
encoding = tokenizer.encode_plus(
    sample_abstract, 
    add_special_tokens=True,
    max_length=128, 
    return_token_type_ids=True, 
    padding="max_length",
    return_attention_mask=True,
    return_tensors="pt"
    )
_, predictions = new_model(encoding["input_ids"], encoding["attention_mask"])
print(predictions)
softmax = nn.Softmax(-1)
print(softmax(predictions))
sample_items["labels"]
df_train.iloc[8][0]<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11022297013041649, 'fraction_numerical': 0.010657691768335437, 'mean_word_length': 5.053904923599321, 'pattern_counts': {'":': 6, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 3, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '97098', 'license': 'MIT', 'n_tokens_mistral': 5010, 'n_tokens_neox': 4902, 'n_words': 885, 'path': 'text classification pytorch lightning.ipynb', 'repo_name': 'kstathou/covid19-classifier', 'size': '14262'}