# ---
# 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
# ---

# + [markdown] id="9tu5aMWDKJon"
# # Verify GPU

# + id="wVbAxbgFccI4" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1619289507794, "user_tz": -480, "elapsed": 589, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="a3921eac-f71e-4a94-9fe8-856993e23f55"
# !nvidia-smi

# + [markdown] id="DBRukcVdFkmc"
# # Weight and Bias (Assisting Metrics, Optional)

# + id="ZJhFk-HiFiMp" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1619289523722, "user_tz": -480, "elapsed": 15278, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="43d837c6-9aad-4d39-fb9d-46e55b15d284"
# !pip install wandb
# !wandb login
project_name = "CoLA with BERT"  # @param {type:"string"}
import os

os.environ["WANDB_PROJECT"] = project_name
import wandb

# + [markdown] id="nv5kop8DFwv9"
# # Installation

# + id="NlWEz_DGdmL2" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1619289531313, "user_tz": -480, "elapsed": 21766, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="66077bf1-64a8-4032-c9ef-37f23d8c07f5"
# !pip install transformers==4.5.1 datasets==1.5.0 accelerate==0.2.1

# + id="NYB87-3cd6VE" executionInfo={"status": "ok", "timestamp": 1619289537813, "user_tz": -480, "elapsed": 12144, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}}
import os

import torch
from accelerate import Accelerator
from datasets import load_dataset, load_metric
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import (
    AdamW,
    BertConfig,
    BertModel,
    BertTokenizerFast,
    default_data_collator,
    get_linear_schedule_with_warmup,
)

# + id="rmPWXk40kfOG" executionInfo={"status": "ok", "timestamp": 1619289537817, "user_tz": -480, "elapsed": 11375, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}}
per_device_train_batch_size = 32  # @param {type:"slider", min:1, max:64, step:1}
learning_rate = 2e-5  # @param {type:"slider", min:2e-7, max:2e-3, step:2e-7}
num_train_epochs = 3  # @param {type:"slider", min:1, max:10, step:1}
logging_steps = 10  # @param {type:"slider", min:10, max:100, step:5}

# + [markdown] id="feEx7HM2KTU3"
# # Tokenizer

# + id="-L9Zno23KTnx" colab={"base_uri": "https://localhost:8080/", "height": 164, "referenced_widgets": ["68d939f55505414c825194e8982054ac", "ff69a9dda1894ff6b88f39a9f81cc2c0", "98d4c9b4ca4f4c00905c3bb228f5605d", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "4e2314037704499c9c16880a1d2e9702", "<KEY>", "434c2be7efaf4a6fa833641fbaf6b11c", "bad1462cc73e4c6ca05625b095db848c", "8091af899341431f9adb3d7c613a3827", "a80f48894f4d4a1fbd24c6c9e5c7a7b0", "<KEY>", "d1488fe68e2d4e58bb134b58858b36e7", "<KEY>", "1983c4310fa848ebb270ad9978acaf30", "5472623fa9ce49a8b7bffc06472d616d", "<KEY>", "<KEY>", "181c118e246440bcb38ce7deeaf87b8b", "<KEY>", "<KEY>", "<KEY>"]} executionInfo={"status": "ok", "timestamp": 1619289540103, "user_tz": -480, "elapsed": 12809, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="578c8448-0ca0-4428-9090-ca055e17205e"
tokenizer = BertTokenizerFast.from_pretrained("bert-base-cased")

# + [markdown] id="3DcNvb7pCbS4"
# # Dataset

# + id="DDJ8tEhBfM_h" colab={"base_uri": "https://localhost:8080/", "height": 218, "referenced_widgets": ["522fccaf1fda4b0eaa9d3cad0b0522a8", "97796a0da4914b83b3a29ee88e58ddf6", "0a145341d4b547279bb4fa89eef2e596", "46a52d4cb17d464583381c60461f9617", "99d9f11b81104b299559717833e91498", "<KEY>", "e054ece1a9944afd819a53ac93241438", "2ba6223403ca4b13a6e282486a5484c6", "<KEY>", "28ca504d4f0e49dabb4357cf7bcc9746", "<KEY>", "769e0622ff8f4174a1f3d96ec3462e37", "684702410ca645e3a5c6f3d7e1d2870d", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "47d1a68f143847aeaf6942d9b74a7941", "<KEY>", "8d22ea682a7044c3b62e12e9dacace47", "<KEY>", "951f1ee32c0e4acebafb8317005fe79b", "669deb9013334792bff7d3647db4e7a1", "<KEY>", "3aadee76307a421a849745859be0deeb", "8f8128e96590401b940be5cd3d8448e1", "<KEY>", "<KEY>", "66217887a9c440f08400a3d252be6e4b", "594d38fd3b29478984ec409e84a5c209", "<KEY>", "be86db13de16402a87e2b384e60e8e95", "<KEY>", "9eb01ee933854cbe88c91841e9d3bf8c", "<KEY>", "6e9661800853459a8e4d64eb1875743e", "91b6e4fe99c44bb1a7ebd4116491a436", "<KEY>", "fa6c21cc679645a0830226708c06463a", "31a3adc322bd408f98af87d703ca518f", "9f534f18938f4b7a8e6fa78dc4a80848", "<KEY>", "8c0509e5ee914eed9e1ccbab742d6726", "724a3d4dd3b0498bacc4e9c022589418", "3393fd4ea83f4b8c9add5704e74b0a7f", "eb348e54e6d84ad78bb8108665715992", "48b25afd11944260be5d4a88fd59f034", "a6a7f95413f448a89598e13a2252f5ec"]} executionInfo={"status": "ok", "timestamp": 1619289542564, "user_tz": -480, "elapsed": 14587, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="0b6b461a-55a4-48a2-d0c1-5c22b5bd1442"
datasets = load_dataset("glue", "cola")
label_list = datasets["train"].features["label"].names


# + id="xGQ8mBlJ5So-" executionInfo={"status": "ok", "timestamp": 1619289542565, "user_tz": -480, "elapsed": 14366, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}}
def preprocess_function(examples):
    # Tokenize the texts
    result = tokenizer(
        examples["sentence"],
        padding="max_length",
        max_length=128,
        truncation="longest_first",
    )
    result["label"] = examples["label"]
    return result


# + id="3mNVI6n65Sk7" colab={"base_uri": "https://localhost:8080/", "height": 164, "referenced_widgets": ["059a77a589644de8a9d3576b08a19a44", "415e955aa86c4c7994bdf81340f99e1b", "1b31010f34154d97949f4fc7715c423e", "d7999a05bc28474fa19aeb875b6c6fc9", "71bfd489e8e94df09d014dc497241c01", "075242e15de2453cbb389249ced66159", "b7029b5aa1104b6c858832b39a8f49a0", "0a11d948faad4e2f8e3ef7881c12c113", "5a9c0200174c414d8c9e5cd6e17c8e7e", "fe375dec6e7844f4a9f9ac4a03383bda", "4415fff623534920b7e96afb1df3e4ef", "<KEY>", "<KEY>", "da661072ad8349ebac756d37f6ff8552", "<KEY>", "<KEY>", "1c635f683d38450db674a54f3e449ff6", "<KEY>", "<KEY>", "25a421036c004f6cac64001c611a7f1a", "<KEY>", "<KEY>", "<KEY>", "ad31a3838aca4e91a53b47d11a3ef08d"]} executionInfo={"status": "ok", "timestamp": 1619289543476, "user_tz": -480, "elapsed": 15027, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="b4a91514-6708-4640-8817-ce7734d47f08"
datasets = datasets.map(preprocess_function, batched=True)
datasets["test"] = datasets["test"].remove_columns("label")

# + id="m0JWvfMW5Si0" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1619289543476, "user_tz": -480, "elapsed": 14734, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="b1bc5f3d-5adb-4377-9ed9-2dd9d1ed5744"
for index in range(3):
    print(f"Sample {index} of the training set: {datasets['train'][index]}.")

# + id="Ow5S4TbKwHJR" executionInfo={"status": "ok", "timestamp": 1619289543477, "user_tz": -480, "elapsed": 14225, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}}
dataloaders = {
    "train": DataLoader(
        datasets["train"],
        collate_fn=default_data_collator,
        batch_size=per_device_train_batch_size,
        shuffle=True,
    ),
    "validation": DataLoader(
        datasets["validation"],
        collate_fn=default_data_collator,
        batch_size=per_device_train_batch_size,
        shuffle=True,
    ),
    "test": DataLoader(
        datasets["test"],
        collate_fn=default_data_collator,
        batch_size=per_device_train_batch_size,
        shuffle=True,
    ),
}


# + [markdown] id="M5_HZzpUCnQY"
# # Model

# + id="cinidhktZubv" executionInfo={"status": "ok", "timestamp": 1619290220230, "user_tz": -480, "elapsed": 914, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}}
class BertClassificationForCoLA(nn.Module):
    def __init__(self, config, from_pretrained=False):
        super().__init__()
        self.config = config
        self.num_labels = config.num_labels

        if not from_pretrained:
            self.bert = BertModel(config)

        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

        self.apply(self._init_weights)  # Equivelent of self.init_weights()

    def _init_weights(self, module):
        """ Initialize the weights """
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    @classmethod
    def from_pretrained(cls, model_name: str, config):
        model = cls(config, from_pretrained=True)
        model.bert = BertModel.from_pretrained(model_name, config=config)
        return model

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        pooled_output = self.dropout(outputs[1])
        logits = self.classifier(pooled_output)

        loss = None
        if labels is not None:
            if self.num_labels == 1:
                #  We are doing regression
                loss_fct = MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

        # return final
        return {
            "loss": loss,
            "logits": logits,
        }


# + id="xvnUY6NJ5SrE" colab={"base_uri": "https://localhost:8080/", "height": 115, "referenced_widgets": ["8154042a7e884162a98e3c81d8781e7b", "<KEY>", "<KEY>", "69a8e2b579a54964916ed6f53dfefa17", "db8597721d1141baa20c10ab79a2a1db", "0f4fa0290f7b442c9a196dc850844ccf", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "cbb4f27d002a42c79246dcf292e204a2", "bafa16537b26423ea395fa98dc862025", "55a8bb11c05b4cc99d6a734840cb479a"]} executionInfo={"status": "ok", "timestamp": 1619289555550, "user_tz": -480, "elapsed": 24721, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="90019c91-d124-43ca-f5ce-f4129539981d"
config = BertConfig.from_pretrained(
    "bert-base-cased", num_labels=len(label_list), finetuning_task="cola"
)
model = BertClassificationForCoLA.from_pretrained("bert-base-cased", config)

# + [markdown] id="eR0zJvATDF4Y"
# # Optimizer, scheduler, accelerator

# + id="6hFGh8Fz8hTU" executionInfo={"status": "ok", "timestamp": 1619289555551, "user_tz": -480, "elapsed": 23640, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}}
no_decay = ["bias", "LayerNorm.weight"]
optimizer_grouped_parameters = [
    {
        "params": [
            p
            for n, p in model.named_parameters()
            if not any(nd in n for nd in no_decay)
        ],
        "weight_decay": 0.01,
    },
    {
        "params": [
            p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)
        ],
        "weight_decay": 0.0,
    },
]
optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate)
num_warmup_steps = 0
scheduler = get_linear_schedule_with_warmup(
    optimizer, num_warmup_steps, num_train_epochs * len(dataloaders["train"])
)

# + id="fs-vxXTKv7zc" executionInfo={"status": "ok", "timestamp": 1619289566142, "user_tz": -480, "elapsed": 33929, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}}
accelerator = Accelerator()
(
    model,
    optimizer,
    dataloaders["train"],
    dataloaders["validation"],
    dataloaders["test"],
) = accelerator.prepare(
    model,
    optimizer,
    dataloaders["train"],
    dataloaders["validation"],
    dataloaders["test"],
)

# + [markdown] id="SASUk8x7DKmO"
# # Training

# + id="0OQKPK5inqrt" colab={"base_uri": "https://localhost:8080/", "height": 567, "referenced_widgets": ["8a8368653bfc4f348f00ee529e4037c4", "687f340b102245dcaf33d7a5aa5fe96b", "65273085867d4f1399a353094d1b9647", "74d73b3b1b424e14a2e1443dea7dceeb", "89ab07e05cc74725875ade04e706f179", "<KEY>", "d7a5408195a84108a799dc375c3ed41c", "05de86eb83094a29a4fa887a5168ea4d", "2bfa1187c0114e949da5d5859acbc351", "<KEY>", "<KEY>", "<KEY>", "e08be46a9f3f45b29bb78e61e73e3aa1", "<KEY>", "2f90e2bd43ef444f85631e34c7d4abe0", "<KEY>", "75f4daca675744bf8a38faf7d0e440e7", "53496062142648a1a2a32a4c4cacfa03", "<KEY>", "f73172225c1f43f1ad60c0be914086c4", "<KEY>", "5ea08c467f0847e6bbe01b0a0254e5f9", "fce1f9b554a946a79c7832e4923fde86", "1852be21114548a78415eb010a3e453f", "<KEY>", "<KEY>", "<KEY>", "<KEY>", "29e268319dbc4e498c7e44c4088d00bf", "90825c16fa7845fdadebfcba0be1c29e", "<KEY>", "7de4fc903af54ba685d116426ae3e913", "<KEY>", "9f17e3c6982e4e54aac6da56b82a05ca", "50b72282e34140d2827d1eef0a9d3efb", "<KEY>", "e2dfe04a610d481194b904c09c87072a", "ea0fffac671848e191a6af7e78ae3df4", "f64de32e5c104311aba65d74cb708e1a", "<KEY>", "7bed6b30fbc74e0db1bb5b849875c68e", "1272fd4bdafe4022884d4eb2889a5e46", "44b7693e922947feab96eb093ee26135", "afa091b45fda4df2aa1ed034d597ee73", "f10bd3c394264a2684f33b636e3f4a6c", "738f7227013e4d51b4dbbabf677b93a2", "0ff2d3eef45e4fb4a450ca9d90559569", "dec07e4ea72a407190134b1927b0fe97", "4266d9baa35448ffa0e5e43900b6ecae", "<KEY>", "f1cc99a685414e568d9bc9eee4ac97ee", "<KEY>", "ed382b34547d42dd9f3e8bb5caaad124", "6a33fdebe534480fba82e8bfcab06d6e", "0c8e283793de4c7b851020d8fa4a3d92", "96f78bff08ba4bc2898d57c50bf55a86"]} executionInfo={"status": "ok", "timestamp": 1619290162516, "user_tz": -480, "elapsed": 582070, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="83962397-dc2d-4914-e1df-df21f503598d"
run_name = "PyTorch Scratch"
if os.getenv("WANDB_PROJECT"):
    wandb.init(project=project_name, name=run_name)
    wandb.watch(model)
for epoch in range(num_train_epochs):
    model.train()
    for step, batch in enumerate(tqdm(dataloaders["train"])):
        batch.pop("idx")

        optimizer.zero_grad()
        outputs = model(**batch)
        loss = outputs["loss"]
        accelerator.backward(loss)
        optimizer.step()
        scheduler.step()

        if step % logging_steps == 0:
            wandb.log(
                {
                    "train/loss": loss,
                    "train/epoch": epoch + step / len(dataloaders["train"]),
                    "train/learning_rate": scheduler.get_last_lr()[0],
                }
            )
    model.eval()
    metric = load_metric("glue", "cola")
    for batch in tqdm(dataloaders["validation"]):
        batch.pop("idx")

        outputs = model(**batch)
        outputs = accelerator.gather(outputs)
        predictions = outputs["logits"].argmax(dim=-1)
        references = accelerator.gather(batch["labels"])

        metric.add_batch(
            predictions=predictions,
            references=references,
        )

    eval_metric = metric.compute()
    wandb.log({"eval/matthews_correlation": eval_metric["matthews_correlation"]})

os.mkdir(run_name)
accelerator.wait_for_everyone()
model = accelerator.unwrap_model(model)
accelerator.save(model.state_dict(), os.path.join(run_name, "trained_model.pt"))

# + [markdown] id="7fa1nAnrDRo1"
# # Test

# + [markdown] id="K6-azaAMdUW4"
# ## On dataset

# + id="m-sYOY_Z5SX2" colab={"base_uri": "https://localhost:8080/", "height": 66, "referenced_widgets": ["ba6ddcf7a3b244b99bf9365c0956ad31", "1f14074fe504477987c0ca85d144439e", "fd29559944534a048a265c519f175e1c", "988cf5f95e174f4996c0f4b8efee3884", "603503008c23465ab76d69c62599ce53", "dd526935d43c4814a9d0d616b536be74", "d547b58b850246e58f82057cae8d9570", "233258573beb4c10946092103b57fb19"]} executionInfo={"status": "ok", "timestamp": 1619290170809, "user_tz": -480, "elapsed": 583369, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="a177bcc9-55f9-4fb8-d848-3c895f172d6b"
predictions = torch.empty(0, dtype=torch.int64, device=accelerator.device)
for batch in tqdm(dataloaders["test"]):
    batch.pop("idx")
    outputs = model(**batch)
    predictions = torch.cat((predictions, outputs["logits"].argmax(dim=-1)), 0)

# + id="KBhIzZc85SVt" tags=[] colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1619290171756, "user_tz": -480, "elapsed": 583954, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="44c059bd-3ac0-410e-81a4-da2f71ab6493"
for index, (sample, pred) in enumerate(zip(datasets["test"]["sentence"], predictions)):
    print(f"{index}\t{label_list[pred.item()]}\t{sample}")

# + [markdown] id="_rfR3-VmdRQn"
# ## Manually

# + id="J2lIJUZSYwmd" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1619290171757, "user_tz": -480, "elapsed": 583280, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="b99b2262-df3d-4387-d91c-47962ea88020"
sentence = "The probable hostile German reaction is unfortunate."  # @param {type:"string"}
tokenized_input = tokenizer(sentence, return_tensors="pt").to(accelerator.device)
outputs = model(**tokenized_input)
print(f"Prediction: {label_list[outputs['logits'].argmax(dim=-1).item()]}")

# + [markdown] id="4D8-YEkriwLX"
# # Inference

# + id="yNJduFfTfahF" executionInfo={"status": "ok", "timestamp": 1619290214182, "user_tz": -480, "elapsed": 808, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}}
import torch
from torch import nn
from transformers import BertConfig, BertModel, BertTokenizerFast

# + id="nPjyt8H_i_8i" executionInfo={"status": "ok", "timestamp": 1619290215730, "user_tz": -480, "elapsed": 2109, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}}
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
label_list = ["unacceptable", "acceptable"]
tokenizer = BertTokenizerFast.from_pretrained("bert-base-cased")
config = BertConfig.from_pretrained("bert-base-cased", finetuning_task="cola")

# + id="M7aVfzCTKdnG" executionInfo={"status": "ok", "timestamp": 1619290230911, "user_tz": -480, "elapsed": 6386, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}}
# Scroll to top and run the class definition class first
model = BertClassificationForCoLA.from_pretrained(
    "PyTorch Scratch/trained_model.pt", config
).to(device)

# + id="PKdUDuHTjlMo" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1619290233772, "user_tz": -480, "elapsed": 641, "user": {"displayName": "<NAME>", "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GhRBTPG7kFx7Gw14fyOvKjff2xqDyZ0tjThv8YKXg=s64", "userId": "13013387071134080776"}} outputId="49e2d96a-2f6c-41ea-f0e8-ebdaff87d364"
sentence = "The probable hostile German reaction is unfortunate."  # @param {type:"string"}
tokenized_input = tokenizer(sentence, return_tensors="pt").to(device)
outputs = model(**tokenized_input)
print(f"Prediction: {label_list[outputs['logits'].argmax(dim=-1).item()]}")

# + id="Natc1mV0KdnG"


StarcoderdataJupyterScriptsDedupFiltered