model.py

import gc

import numpy as np
import pandas as pd

import torch
from torch import nn
import transformers
from transformers import AutoModel, AutoTokenizer, AutoConfig


config = dict(
    # basic
    seed = 3407,
    num_jobs=1,
    num_labels=2,
    
    # model info
    tokenizer_path = 'roberta-large', # 'allenai/biomed_roberta_base',
    model_checkpoint = 'roberta-large', # 'allenai/biomed_roberta_base', 
    device = 'cuda' if torch.cuda.is_available() else 'cpu',
    
    # training paramters
    max_length = 512,
    batch_size=16,
    
    # for this notebook
    debug = False,
)


def create_sample_test():
    feats = pd.read_csv(f"../input/nbme-score-clinical-patient-notes/features.csv")
    feats.loc[27, 'feature_text'] = "Last-Pap-smear-1-year-ago"
    
    notes = pd.read_csv(f"../input/nbme-score-clinical-patient-notes/patient_notes.csv")
    test = pd.read_csv(f"../input/nbme-score-clinical-patient-notes/test.csv")

    merged = test.merge(notes, how = "left")
    merged = merged.merge(feats, how = "left")

    def process_feature_text(text):
        return text.replace("-OR-", ";-").replace("-", " ")
    merged["feature_text"] = [process_feature_text(x) for x in merged["feature_text"]]
    
    return merged.sample(1).reset_index(drop=True)

class NBMETestData(torch.utils.data.Dataset):
    def __init__(self, feature_text, pn_history, tokenizer):
        self.feature_text = feature_text
        self.pn_history = pn_history
        self.tokenizer = tokenizer
    
    def __len__(self):
        return len(self.feature_text)
    
    def __getitem__(self, idx):
        tokenized = self.tokenizer(
            self.feature_text[idx],
            self.pn_history[idx],
            truncation = "only_second",
            max_length = config['max_length'],
            padding = "max_length",
            return_offsets_mapping = True
        )
        tokenized["sequence_ids"] = tokenized.sequence_ids()

        input_ids = np.array(tokenized["input_ids"])
        attention_mask = np.array(tokenized["attention_mask"])
        offset_mapping = np.array(tokenized["offset_mapping"])
        sequence_ids = np.array(tokenized["sequence_ids"]).astype("float16")

        return {
            'input_ids': input_ids, 
            'attention_mask': attention_mask, 
            'offset_mapping': offset_mapping, 
            'sequence_ids': sequence_ids,
        }

# class NBMEModel(nn.Module):
#     def __init__(self, num_labels=1, path=None):
#         super().__init__()
        
#         layer_norm_eps: float = 1e-6
        
#         self.path = path
#         self.num_labels = num_labels
        
#         self.transformer = transformers.AutoModel.from_pretrained(config['model_checkpoint'])
#         self.dropout = nn.Dropout(0.2)
#         self.output = nn.Linear(768, 1)
        
#         if self.path is not None:
#             self.load_state_dict(torch.load(self.path)['model'])

#     def forward(self, data):
        
#         ids = data['input_ids']
#         mask = data['attention_mask']
#         try:
#             target = data['targets']
#         except:
#             target = None

#         transformer_out = self.transformer(ids, mask)
#         sequence_output = transformer_out[0]
#         sequence_output = self.dropout(sequence_output)
#         logits = self.output(sequence_output)

#         ret = {
#             "logits": torch.sigmoid(logits),
#         }
        
#         if target is not None:
#             loss = self.get_loss(logits, target)
#             ret['loss'] = loss
#             ret['targets'] = target

#         return ret

        
#     def get_optimizer(self, learning_rate, weigth_decay):
#         optimizer = torch.optim.AdamW(
#             self.parameters(), 
#             lr=learning_rate, 
#             weight_decay=weigth_decay,
#         )
#         if self.path is not None:
#             optimizer.load_state_dict(torch.load(self.path)['optimizer'])
        
#         return optimizer
            
#     def get_scheduler(self, optimizer, num_warmup_steps, num_training_steps):
#         scheduler = transformers.get_linear_schedule_with_warmup(
#             optimizer,
#             num_warmup_steps=num_warmup_steps,
#             num_training_steps=num_training_steps,
#         )
#         if self.path is not None:
#             scheduler.load_state_dict(torch.load(self.path)['scheduler'])
            
#         return scheduler
    
#     def get_loss(self, output, target):
#         loss_fn = nn.BCEWithLogitsLoss(reduction="none")
#         loss = loss_fn(output.view(-1, 1), target.view(-1, 1))
#         loss = torch.masked_select(loss, target.view(-1, 1) != -100).mean()
#         return loss


class NBMEModel(nn.Module):
    def __init__(self, num_labels=2, path=None):
        super().__init__()
        
        layer_norm_eps: float = 1e-6
        
        self.path = path
        self.num_labels = num_labels
        self.transformer = transformers.AutoModel.from_pretrained(config['model_checkpoint'])
        self.dropout = nn.Dropout(0.1)
        
        self.dropout1 = nn.Dropout(0.1)
        self.dropout2 = nn.Dropout(0.2)
        self.dropout3 = nn.Dropout(0.3)
        self.dropout4 = nn.Dropout(0.4)
        self.dropout5 = nn.Dropout(0.5)
        
        self.output = nn.Linear(1024, 1)
        
        if self.path is not None:
            self.load_state_dict(torch.load(self.path)['model'])
    
    def forward(self, data):
        
        ids = data['input_ids']
        mask = data['attention_mask']
        try:
            target = data['targets']
        except:
            target = None

        transformer_out = self.transformer(ids, mask)
        sequence_output = transformer_out[0]
        sequence_output = self.dropout(sequence_output)
    
        logits1 = self.output(self.dropout1(sequence_output))
        logits2 = self.output(self.dropout2(sequence_output))
        logits3 = self.output(self.dropout3(sequence_output))
        logits4 = self.output(self.dropout4(sequence_output))
        logits5 = self.output(self.dropout5(sequence_output))

        logits = (logits1 + logits2 + logits3 + logits4 + logits5) / 5
        ret = {
            'logits': torch.sigmoid(logits), 
        }
        
        loss = 0

        if target is not None:
            loss1 = self.get_loss(logits1, target)
            loss2 = self.get_loss(logits2, target)
            loss3 = self.get_loss(logits3, target)
            loss4 = self.get_loss(logits4, target)
            loss5 = self.get_loss(logits5, target)
            loss = (loss1 + loss2 + loss3 + loss4 + loss5) / 5
            ret['loss'] = loss
            ret['target'] = target

        return ret

        
    def get_optimizer(self, learning_rate, weigth_decay):
        optimizer = torch.optim.AdamW(
            self.parameters(), 
            lr=learning_rate, 
            weight_decay=weigth_decay,
        )
        if self.path is not None:
            optimizer.load_state_dict(torch.load(self.path)['optimizer'])
        
        return optimizer
            
    def get_scheduler(self, optimizer, num_warmup_steps, num_training_steps):
        scheduler = transformers.get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=num_warmup_steps,
            num_training_steps=num_training_steps,
        )
        if self.path is not None:
            scheduler.load_state_dict(torch.load(self.path)['scheduler'])
            
        return scheduler
    
    def get_loss(self, output, target):
        loss_fn = nn.BCEWithLogitsLoss(reduction="none")
        loss = loss_fn(output.view(-1, 1), target.view(-1, 1))
        loss = torch.masked_select(loss, target.view(-1, 1) != -100).mean()
        return loss


def get_location_predictions(preds, offset_mapping, sequence_ids, test=False):
    all_predictions = []
    for pred, offsets, seq_ids in zip(preds, offset_mapping, sequence_ids):
        start_idx = None
        current_preds = []
        for p, o, s_id in zip(pred, offsets, seq_ids):
            if s_id is None or s_id == 0:
                continue
            if p > 0.5:
                if start_idx is None:
                    start_idx = o[0]
                end_idx = o[1]
            elif start_idx is not None:
                if test:
                    current_preds.append(f"{start_idx} {end_idx}")
                else:
                    current_preds.append((start_idx, end_idx))
                start_idx = None
        if test:
            all_predictions.append("; ".join(current_preds))
        else:
            all_predictions.append(current_preds)
    return all_predictions



def predict_location_preds(tokenizer, model, feature_text, pn_history, pn_history_lower):

    test_ds = NBMETestData(feature_text, pn_history_lower, tokenizer)
    test_dl = torch.utils.data.DataLoader(
        test_ds, 
        batch_size=config['batch_size'], 
        pin_memory=True, 
        shuffle=False, 
        drop_last=False
    )

    all_preds = None
    offsets = []
    seq_ids = []

    preds = []

    with torch.no_grad():
        for batch in test_dl:

            for k, v in batch.items():
                if k not in  ['offset_mapping', 'sequence_id']:
                    batch[k] = v.to(config['device'])

            logits = model(batch)['logits']
            preds.append(logits.cpu().numpy())

            offset_mapping = batch['offset_mapping']
            sequence_ids = batch['sequence_ids']
            offsets.append(offset_mapping.cpu().numpy())
            seq_ids.append(sequence_ids.cpu().numpy())

    preds = np.concatenate(preds, axis=0)
    if all_preds is None:
        all_preds = np.array(preds).astype(np.float32)
    else:
        all_preds += np.array(preds).astype(np.float32)
    torch.cuda.empty_cache()

    all_preds = all_preds.squeeze()

    offsets = np.concatenate(offsets, axis=0)
    seq_ids = np.concatenate(seq_ids, axis=0)

    # print(all_preds.shape, offsets.shape, seq_ids.shape)

    location_preds = get_location_predictions([all_preds], offsets, seq_ids, test=False)[0]
    
    x = []
    
    for location in location_preds:
        x.append(pn_history[0][location[0]: location[1]])
    
    return location_preds, ', '.join(x)

def get_predictions(feature_text, pn_history):
    feature_text = feature_text.lower().replace("-OR-", ";-").replace("-", " ")
    pn_history_lower = pn_history.lower()
    
    location_preds, pred_string = predict_location_preds(tokenizer, model, [feature_text], [pn_history], [pn_history_lower])
    
    if pred_string == "":
        pred_string = 'Feature not present!'
    else:
        pred_string = 'Feature is present!' + '\nText Span - ' + pred_string
    
    return pred_string

tokenizer = AutoTokenizer.from_pretrained(config['tokenizer_path'])
path = 'model_large_pseudo_label.pth'

model = NBMEModel().to(config['device'])
model.load_state_dict(
    torch.load(
        path, 
        map_location=torch.device(config['device'])
    )
)
model.eval()