run_bionlp_search.py

from torch.optim import AdamW
from torch.utils.data import DataLoader
from datasets import load_dataset
from transformers import (
    AutoModelForTokenClassification,
    AutoTokenizer,
    DataCollatorForTokenClassification,
    TrainingArguments,
    Trainer,
    get_linear_schedule_with_warmup,
)

from bionlp import compute_metrics,label2id,id2label,label_list

import torch
import numpy as np

from betty.engine import Engine
from betty.configs import Config, EngineConfig
from betty.problems import ImplicitProblem


class Architecture(torch.nn.Module):
    def __init__(self, r_search, n_layer):
        super(Architecture, self).__init__()
        self.alphas=torch.nn.Parameter(torch.zeros(n_layer,r_search*2))

    def forward(self):
        return self.alphas

class Arch(ImplicitProblem):
    def training_step(self, batch):

        alphas = self.forward()
        #print('test:arch')
        loss = self.roberta.module(alphas, **batch)  
        #print('test:arch2',loss)
        return loss


class Roberta(ImplicitProblem):
    def training_step(self, batch):

        alphas = self.arch()
        #print('test:gpt')
        #print(batch)

        loss = self.module(alphas, **batch).loss 
 
        #print('test:gpt2',loss)
        return loss


class NASEngine(Engine):
    @torch.no_grad()
    def validation(self):
        #print('test:validation')
        alphas_q,alphas_v=self.arch.module.alphas[:,:lora_dim],self.arch.module.alphas[:,lora_dim:]
        r_list_q=[int(sum(torch.nn.functional.softmax(alpha,dim=-1)>=1/lora_dim)) for alpha in alphas_q]
        r_list_v=[int(sum(torch.nn.functional.softmax(alpha,dim=-1)>=1/lora_dim)) for alpha in alphas_v]

        return {"q": str(r_list_q),'v':str(r_list_v)}
    
model_checkpoint = "roberta-base"
num_layers=12
lora_dim=8
lora_alpha=16
lr = 1e-3
batch_size = 16
num_epochs = 50
device='cuda'

bionlp = load_dataset("tner/bionlp2004")


tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, add_prefix_space=True)

def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
    labels = []
    for i, label in enumerate(examples[f"tags"]):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx is None:
                label_ids.append(-100)
            elif word_idx != previous_word_idx:
                label_ids.append(label[word_idx])
            else:
                label_ids.append(-100)
            previous_word_idx = word_idx
        labels.append(label_ids)
    tokenized_inputs["labels"] = labels
    return tokenized_inputs

tokenized_bionlp = bionlp.map(tokenize_and_align_labels, batched=True)

#print(tokenized_bionlp['train'].remove_columns("tokens")[0])
#exit()

data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)
def collate_fn(examples):
    return tokenizer.pad(examples, padding="longest", return_tensors="pt",)


train_dataset=tokenized_bionlp['train'].train_test_split(test_size=0.3)
train_dataset=train_dataset.remove_columns(["tokens",'tags'])
train_split=train_dataset['train']
search_split=train_dataset['test']
# Instantiate dataloaders.
train_dataloader = DataLoader(train_split, shuffle=True, collate_fn=data_collator, batch_size=batch_size,drop_last=True)
search_dataloader = DataLoader(search_split, shuffle=True, collate_fn=data_collator, batch_size=batch_size,drop_last=True)


model = AutoModelForTokenClassification.from_pretrained(
    model_checkpoint, num_labels=11, id2label=id2label, label2id=label2id,
    apply_lora=True,
    lora_alpha=lora_alpha,
    lora_r=lora_dim,
)

lora_parameters =[
        {
            "params": [p for n, p in model.named_parameters() if 'lora_' in n], # if not any(nd in n for nd in no_decay)],
        }]
optimizer = AdamW(params=lora_parameters, lr=lr)
arch_net=Architecture(lora_dim,num_layers)

arch_optimizer = torch.optim.Adam(
    arch_net.parameters(),
    lr=3e-4,
    betas=(0.5, 0.999),
    weight_decay=1e-3,
)
# Instantiate scheduler
lr_scheduler = get_linear_schedule_with_warmup(
    optimizer=optimizer,
    num_warmup_steps=0.06 * (len(train_dataloader) * num_epochs),
    num_training_steps=(len(train_dataloader) * num_epochs),
)

outer_config = Config(retain_graph=True)
unroll_steps=1
inner_config = Config(type="darts", unroll_steps=unroll_steps)

outer = Arch(
    name="arch",
    module=arch_net,
    optimizer=arch_optimizer,
    train_data_loader=search_dataloader,
    config=outer_config,
)
inner = Roberta(
    name="roberta",
    module=model,
    optimizer=optimizer,
    train_data_loader=train_dataloader,
    config=inner_config,
)

problems = [outer, inner]
l2u = {inner: [outer]}
u2l = {outer: [inner]}
dependencies = {"l2u": l2u, "u2l": u2l}


train_portion=1.0
report_freq = 100

num_train = len(train_split)  # 50000
indices = list(range(num_train))
split = int(np.floor(train_portion * num_train))

train_iters = int(
    num_epochs
    * (num_train * train_portion // batch_size + 1)
    * unroll_steps
)
print(train_iters,num_train)

engine_config = EngineConfig(
    valid_step=report_freq * unroll_steps,
    train_iters=train_iters,
    roll_back=True,
)

engine = NASEngine(config=engine_config, problems=problems, dependencies=dependencies)
engine.run()

'''
for epoch in range(num_epochs):
    model.train()
    #for step, batch in enumerate(tqdm(train_dataloader)):
    epoch_loss,epoch_loss_search=0,0
    for step, batch in enumerate((train_dataloader)):
        batch.to(device)
        outputs = model(**batch)
        loss = outputs.loss
        epoch_loss += batch_size * loss.item()
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        if globals.search:
            batch = next(iter(search_dataloader))
            #print(batch)
            batch.to(device)
            outputs = model(**batch)
            loss = outputs.loss
            epoch_loss_search += batch_size * loss.item()
            loss.backward()
            optimizer_a.step()
            optimizer_a.zero_grad()
        #print(model.base_model.model.roberta.encoder.layer[0].attention.self.query.a)
    print('train loss epoch ',epoch+1, epoch_loss / len(train_split))

    if globals.search:
        print('train loss search epoch ',epoch+1, epoch_loss_search / len(train_split))
        if model_checkpoint == "roberta-base" or model_checkpoint == "roberta-large":
            query_rlist=[int(sum(torch.nn.functional.softmax(layer.attention.self.query.a,dim=0)>=1/lora_dim)) for layer in layers]
            value_rlist=[int(sum(torch.nn.functional.softmax(layer.attention.self.value.a,dim=0)>=1/lora_dim)) for layer in layers]
        else:
            query_rlist=[int(sum(torch.nn.functional.softmax(layer.attention.self.query_proj.a,dim=0)>=1/lora_dim)) for layer in layers]
            value_rlist=[int(sum(torch.nn.functional.softmax(layer.attention.self.value_proj.a,dim=0)>=1/lora_dim)) for layer in layers]
        rlist=[]
        for i in range(num_layers):
            rlist.append(query_rlist[i])
            rlist.append(value_rlist[i])
        
        print(rlist)
    else:
        model.eval()
        pred,ref=[],[]
        #for step, batch in enumerate(tqdm(eval_dataloader)):
        for step, batch in enumerate((eval_dataloader)):
            batch.to(device)
            with torch.no_grad():
                outputs = model(**batch)
            predictions = outputs.logits.argmax(dim=-1)
            pred.append(predictions.cpu().detach().numpy())
            ref.append(batch['labels'].cpu().detach().numpy())
            #predictions, references = predictions, batch["labels"]
            #metric.add_batch(predictions=predictions,references=references)

        eval_metric = compute_metrics(zip(pred,ref))
        print(f"epoch {epoch}:", eval_metric)
'''