# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + colab={"base_uri": "https://localhost:8080/"} id="XgP4Coqt2on4" outputId="f26b8027-f1df-4cb2-dee7-c3d40ece948a"
import os
os.chdir('/content/')
!(echo 'Clone LongSeq git ...'; git clone https://github.com/AUTH-MINT/LongSeq.git; echo 'done')
# !test -d /content/ebm_nlp/ || (echo 'Download EBM-NLP dataset'; git clone https://github.com/bepnye/EBM-NLP.git ebm_nlp; echo 'done')

os.chdir('/content/LongSeq/')
# !test -d /content/LongSeq/ebm_nlp_2_00/ || (echo 'Extracting documents and labels...'; tar -xzf /content/ebm_nlp/ebm_nlp_2_00.tar.gz; echo 'done')

os.chdir('/content/')
!(echo 'Delete unnecessary files...'; rm -rf /content/ebm_nlp/; echo 'done')
print('All ready!')

# + colab={"base_uri": "https://localhost:8080/"} id="FTiTfyk127Ih" outputId="52adea95-39ba-4e5b-870c-d11023512d82"
print("Installing transformers...\n")
# !pip install transformers==3.3.1
print("\n")
print("Installing pytorch crf...\n")
# !pip install pytorch-crf
print("\n")

# + colab={"base_uri": "https://localhost:8080/", "height": 113, "referenced_widgets": ["<KEY>", "87b358e5b4d948f68b95380a20d5e340", "eb657b6d4f64426d9e7ca9d4e6d07ff1", "3a411d7fc59847e79bd1ab92ca447ca1", "<KEY>", "<KEY>", "<KEY>", "9e696541c0a04b46921469f4895997eb", "3c4318a541634e18833f6e46451a6b88", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "c0384becebdc4160bf0a75c6b5b3fcc5", "4621e307021441fd85de6c13742508e1", "29b92e3403fc42798064b82ddb35966e", "fc5210457e7f417490bbe432fa66416a", "88c2811687e44c838e0a06598068607c", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "1233e60d7f98416bafba6f5ed95d9c0c", "06fcc918ed58482f98f8b3cd1bb1940e", "cc5ede8891ea4fa5b55823f334c0f7cc", "ab41a2f99e374a6d90f2c8183099e715", "<KEY>", "<KEY>", "ba4c3e5718cb4b9eb395196274fde0b3"]} id="0GdipjP36nqK" outputId="e24894c4-0c4b-4b70-8b6a-9c1bf52209ed"
os.chdir('/content/LongSeq/')
import torch
import numpy as np
import os
#import nltk
import build_data
import build_data_abstracts
from transformers import AdamW
from transformers import get_linear_schedule_with_warmup
import time
import model.dataset as dataset
import model.engine as engine
from model.modeltr import EntityModel
from model.data_utils import CoNLLDataset, data_to_dict
from model.config import Config
from model.utils import timeSince, metrics
import model.utils as utils

# + colab={"base_uri": "https://localhost:8080/"} id="zAJtvjky6-ea" outputId="6327dc67-104e-4942-927f-642b3b7d54c5"
config = Config()
 
data_prefix = 'p1_all'
 
print("Check if files exist otwerwise create them!")

print("Use sentences:",config.use_sen)
 
if config.use_sen:
    print("Dataset Contains each sentence.")
    print('Generating tagged and formatted data for the model')
    build_data.generate_model_data(data_prefix)
    print('Finished. Ready to train model on %s' %data_prefix)
    cwd = os.getcwd()
    config.filename_dev = os.path.join(cwd, 'data', data_prefix + '_' + os.path.basename(config.filename_dev))
    config.filename_test = os.path.join(cwd, 'data', data_prefix + '_' + os.path.basename(config.filename_test))
    config.filename_train = os.path.join(cwd, 'data', data_prefix + '_' + os.path.basename(config.filename_train))
    in_tp = 'Sentences'
 
else:
    print("Dataset Contains whole abstract.")
    print('Generating tagged and formatted data for the model')
    build_data_abstracts.generate_model_data(data_prefix)#, voc_build=False)
    print('Finished. Ready to train model on %s' % data_prefix)
    cwd = os.getcwd()
    config.filename_dev = os.path.join(cwd, 'data', data_prefix + '_' + os.path.basename(config.filename_dev_abs))
    config.filename_test = os.path.join(cwd, 'data', data_prefix + '_' + os.path.basename(config.filename_test_abs))
    config.filename_train = os.path.join(cwd, 'data', data_prefix + '_' + os.path.basename(config.filename_train_abs))
    in_tp = "Abstract"
 
print("All done...")

# + colab={"base_uri": "https://localhost:8080/"} id="3VUUCYJo7ON9" outputId="33775662-7014-4c7a-d852-da7b93d8a4a3"
parameters = {
    'Use:':in_tp,
    'Pretrained Language Model:':Config.PRETRAINED_MODEL,
    'Number of epochs:':Config.nepochs,
    'Learning Rate:':Config.lr,
    'N:':Config.N,
    'h:':Config.h,
    'd_model:':Config.d_model,
    'Train_Batch_Size':Config.train_batch_size,
    'Test_Batch_Size':Config.valid_batch_size
}

print("Parameters:\n")
for i in parameters:
    print(i,parameters[i])


# + colab={"base_uri": "https://localhost:8080/"} id="urkZ3IGoDf0o" outputId="03e007ff-76a7-4c59-dc5a-a03015980517"
# create datasets
print("Load datasets:")

dev = CoNLLDataset(config.filename_dev, config.max_iter)
      
train = CoNLLDataset(config.filename_train, config.max_iter)
      
test  = CoNLLDataset(config.filename_test, config.max_iter)
      

# + colab={"base_uri": "https://localhost:8080/"} id="uFHrRgeuDgh4" outputId="095a51fe-c164-480f-e9d8-fac2b2b515c7"
print('Transform data to dictionaries:')
devDict = data_to_dict(dev)
trainDict = data_to_dict(train)
testDict = data_to_dict(test)
num_tag = len(list(utils.enc_trans.classes_))

# + colab={"base_uri": "https://localhost:8080/"} id="fR6L2nHKDi3p" outputId="5ac08d66-3561-4f3a-ee92-5c848d69a41a"
print(devDict['words'][0])

# + colab={"base_uri": "https://localhost:8080/"} id="_4K_1vrsDl14" outputId="2d27c55e-a226-4928-bff7-59c9952ffc93"
print('Initialize dataloaders:')

train_dataset = dataset.EntityDataset(
        words=trainDict['words'], tags=trainDict['enc_tags'], config=config
    )

train_data_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=Config.train_batch_size, num_workers=4
    )

valid_dataset = dataset.EntityDataset(
        words=devDict['words'], tags=devDict['enc_tags'], config=config
    )

valid_data_loader = torch.utils.data.DataLoader(
        valid_dataset, batch_size=Config.valid_batch_size, num_workers=1
    )

test_dataset = dataset.EntityDataset(
        words=testDict['words'], tags=testDict['enc_tags'], config=config
    )

test_data_loader = torch.utils.data.DataLoader(
        test_dataset, batch_size=1, num_workers=1
    )


# + colab={"base_uri": "https://localhost:8080/"} id="7GLiVS1-Do-5" outputId="7c14be0f-0b2b-4e69-d474-09f605ab4775"
device = torch.device("cuda")
model = EntityModel(num_tag=num_tag, tr=True)
model.to(device)

# + id="9_kgbB1GDq-I"
param_optimizer = list(model.named_parameters())
no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
optimizer_parameters = [
        {
            "params": [
                p for n, p in param_optimizer if not any(nd in n for nd in no_decay)
            ],
            "weight_decay": 0.001,
        },
        {
            "params": [
                p for n, p in param_optimizer if any(nd in n for nd in no_decay)
            ],
            "weight_decay": 0.0,
        },
    ]

# + id="3vQuaHxID6Wt"
num_train_steps = int(len(trainDict['words']) / Config.train_batch_size * Config.nepochs)
optimizer = AdamW(optimizer_parameters, lr=Config.lr)
scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=0, num_training_steps=num_train_steps)

# + colab={"base_uri": "https://localhost:8080/"} id="MCNgtI0jD8Vx" outputId="bf0e3d7d-efbf-4de1-ea8f-a34d1aaa9e4a"
print("Training starts...\n\n")
start = time.time()

best_loss = np.inf
for epoch in range(Config.nepochs):
    print("Epoch %g from %g"%(epoch+1,Config.nepochs))
    train_loss = engine.train_fn(train_data_loader, model, optimizer, device, scheduler)
    test_loss = engine.eval_fn(valid_data_loader, model, device)
    print(f"Train Loss = {train_loss} Valid Loss = {test_loss}")
    if test_loss < best_loss:
        torch.save(model.state_dict(), Config.MODEL_PATH)
        best_loss = test_loss
        print("\nNew best score!\n")

total_time = timeSince(start)

print('Total training time:',total_time)
parameters["Training Time:"] = total_time

# + colab={"base_uri": "https://localhost:8080/"} id="SnF8_ES7D-aB" outputId="94b4377f-3015-4b0e-aab9-29edba6fe2b2"
print("Testing...\n")

model = EntityModel(num_tag=num_tag, tr=False)
model.load_state_dict(torch.load(Config.MODEL_PATH))
model.to(device)
results = engine.test_fn(test_data_loader, model, device)

y = results['y']
y_pred = results['y_pred']


print("Labels, no X")
res3 = metrics(y,y_pred, labels=utils.def_tags[:-1], target_names=utils.labels)
print("Parameters:\n")
for i in parameters:
    print(i,parameters[i])

# + colab={"base_uri": "https://localhost:8080/", "height": 34} id="ZXTOgGoVeZF-" outputId="49cc86fe-6969-421d-bed9-598a592df7b6"
from google.colab import files
files.download("/content/res.zip")

# + id="sJ9la6CHfM0a"


StarcoderdataJupyterScriptsDedupFiltered