# ---
# 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="lIYdn1woOS1n" outputId="cece5524-0d94-4cc4-b260-23e2f0ecc744"


# + id="Y2upWg_Qvax1"
import functools

import datasets

import torchtext
import torch
import torch.nn as nn
import torch.optim as optim

# + colab={"base_uri": "https://localhost:8080/"} id="ZIVeVqVUvdcK" outputId="db0dbf36-4a75-4d30-bcef-8cb52b7a5b30"
train_data, test_data = datasets.load_dataset('imdb', split=['train', 'test'])

# + colab={"base_uri": "https://localhost:8080/"} id="2f6USIQOvkV_" outputId="3444d1ad-e2eb-4766-92d8-ea4d85f7011f"
train_data, test_data

# + id="3Av1UUtZxoSL"
tokenizer = torchtext.data.utils.get_tokenizer('basic_english')


# + id="ju4G0eKEx2RO"
def tokenize_data(example, tokenizer):
    tokens = {'tokens': tokenizer(example['text'])}
    return tokens


# + colab={"base_uri": "https://localhost:8080/"} id="T6RecxI4xs7s" outputId="0c97c5cb-7bca-4664-98fa-5b665ab510c9"
train_data = train_data.map(tokenize_data, fn_kwargs={'tokenizer': tokenizer})
test_data = test_data.map(tokenize_data, fn_kwargs={'tokenizer': tokenizer})

# + colab={"base_uri": "https://localhost:8080/"} id="-Foj4qesxqiz" outputId="09fc27eb-0f50-45df-aacb-debf53948cdc"
train_data, test_data

# + id="OfzYwhN2wFcA"
train_valid_data = train_data.train_test_split(test_size=0.25)
train_data = train_valid_data['train']
valid_data = train_valid_data['test']

# + colab={"base_uri": "https://localhost:8080/"} id="ovGkrOJkwZsC" outputId="b5ca62c6-96e5-4eca-8c0f-558fd087dee1"
len(train_data), len(valid_data), len(test_data)

# + id="9VMeQG_FxUVO"
min_freq = 3
special_tokens = ['<unk>', '<pad>']

vocab = torchtext.vocab.build_vocab_from_iterator(train_data['tokens'],
                                                  min_freq=min_freq,
                                                  specials=special_tokens)

# + colab={"base_uri": "https://localhost:8080/"} id="rbroBAClxXGB" outputId="91c5da92-7f97-4ad8-a946-6de1899e64a2"
len(vocab)

# + colab={"base_uri": "https://localhost:8080/"} id="3bKHqCxPyQSb" outputId="35b3c437-f0f8-43fb-8f10-b968c597d5b4"
vocab.get_itos()[:10]

# + colab={"base_uri": "https://localhost:8080/"} id="uStvd2szyUGR" outputId="d4c9d2a4-86a9-413f-9200-5f1a27ece925"
unk_index = vocab['<unk>']

unk_index

# + colab={"base_uri": "https://localhost:8080/"} id="gd5R8NCJyws4" outputId="46666a6d-56ff-42e3-ebe2-6f87930270c7"
pad_index = vocab['<pad>']

pad_index

# + id="_syj_YR8yp7B"
vocab.set_default_index(unk_index)


# + id="ENlE1eAM0lHe"
def numericalize_data(example, vocab):
    ids = {'ids': [vocab[token] for token in example['tokens']]}
    return ids


# + colab={"base_uri": "https://localhost:8080/", "height": 164, "referenced_widgets": ["ab96a5a035b140aa9958ada5eeaa8edc", "e393253753984998b533c18502dd4477", "9c6b3d3c5c5a4c0b940c0aec3d833f58", "a24b5ac5e18b48d48467217c8ea67463", "b699f48ad80442d38481d4f960cf23a3", "707426cd59454f53b955a2ed9e46b90e", "<KEY>", "38e9f642f0714b95ba51da7cafbd1ff2", "852f269e78204293afba0e452e43eba2", "c87dea49f30b471d9128b2b2082adf89", "<KEY>", "6c44fd275fa34a2ab02b16519e48aca7", "<KEY>", "398dcca3803c46beb2854bd52d66e322", "<KEY>", "0cf052dd17f54df19b0e5de1891ec1a4", "<KEY>", "<KEY>", "<KEY>", "d52ef97893144dffa64d1da2b8ea52e6", "d0832872ea074689b8c7e3b8b647e68d", "<KEY>", "eeac1f30dad64508ba7f1a9473f390bd", "958e2685a018455b9ea02803ca96b464"]} id="ux_YLzDA069-" outputId="11d67399-ace5-49ed-8327-0f85de3386a6"
train_data = train_data.map(numericalize_data, fn_kwargs={'vocab': vocab})
valid_data = valid_data.map(numericalize_data, fn_kwargs={'vocab': vocab})
test_data = test_data.map(numericalize_data, fn_kwargs={'vocab': vocab})

# + id="GAXGqlXT1FD9"
train_data.set_format(type='torch', columns=['ids', 'label'])
valid_data.set_format(type='torch', columns=['ids', 'label'])
test_data.set_format(type='torch', columns=['ids', 'label'])


# + id="qCmyoFKAvmnj"
class NBoW(nn.Module):
    def __init__(self, vocab_size, embedding_dim, output_dim, pad_index):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=pad_index)
        self.fc = nn.Linear(embedding_dim, output_dim)

    def forward(self, text):
        # text = [batch size, seq len]
        embedded = self.embedding(text)
        # embedded = [batch size, seq len, embedding dim]
        pooled = embedded.mean(dim=1)
        # pooled = [batch size, embedding dim]
        prediction = self.fc(pooled)
        # prediction = [batch size, output dim]
        return prediction


# + id="lPym0qxrwC7e"
vocab_size = len(vocab)
embedding_dim = 256
output_dim = 2

model = NBoW(vocab_size, embedding_dim, output_dim, pad_index)

# + colab={"base_uri": "https://localhost:8080/"} id="Prvx6C3TFyI4" outputId="c9069fbe-ec6e-423d-ff59-820f934c6432"
vectors = torchtext.vocab.FastText()

# + colab={"base_uri": "https://localhost:8080/"} id="e3gI343FIETN" outputId="44ed85d4-3e22-4bf0-e86e-0d7c01e33c35"
vectors = torchtext.vocab.GloVe()

# + id="xLg8TKFCzAfL"
optimizer = optim.Adam(model.parameters())

# + id="XJrIlwlfzQqY"
criterion = nn.CrossEntropyLoss()

# + colab={"base_uri": "https://localhost:8080/"} id="KbEREGWmzR7J" outputId="44ec5013-ab39-4f35-b425-250f7e4093ef"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

device

# + id="IBtd-0I3zVTo"
model = model.to(device)
criterion = criterion.to(device)


# + id="TNaDUz3M2QDM"
def collate(batch, pad_index):
    batch_ids = [i['ids'] for i in batch]
    batch_ids = nn.utils.rnn.pad_sequence(batch_ids, padding_value=pad_index, batch_first=True)
    batch_labels = [i['label'] for i in batch]
    batch_labels = torch.stack(batch_labels)
    batch = {'ids': batch_ids,
             'labels': batch_labels}
    return batch


# + id="LYsAzjrV0AnH"
batch_size = 512

collate = functools.partial(collate, pad_index=pad_index)

train_dataloader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, collate_fn=collate)
valid_dataloader = torch.utils.data.DataLoader(valid_data, batch_size=batch_size, collate_fn=collate)
test_dataloader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, collate_fn=collate)


# + id="mKeLtjK5zX41"
def train(dataloader, model, criterion, optimizer, device):

    model.train()
    epoch_loss = 0
    epoch_accuracy = 0

    for batch in dataloader:
        tokens = batch['ids'].to(device)
        labels = batch['labels'].to(device)
        predictions = model(tokens)
        loss = criterion(predictions, labels)
        accuracy = get_accuracy(predictions, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        epoch_accuracy += accuracy.item()

    return epoch_loss / len(dataloader), epoch_accuracy / len(dataloader)


# + id="3gJHwUZZ0NC6"
def evaluate(dataloader, model, criterion, device):
    
    model.eval()
    epoch_loss = 0
    epoch_accuracy = 0

    with torch.no_grad():
        for batch in dataloader:
            tokens = batch['ids'].to(device)
            labels = batch['labels'].to(device)
            predictions = model(tokens)
            loss = criterion(predictions, labels)
            accuracy = get_accuracy(predictions, labels)
            epoch_loss += loss.item()
            epoch_accuracy += accuracy.item()

    return epoch_loss / len(dataloader), epoch_accuracy / len(dataloader)


# + id="DOPRg4Gg5L44"
def get_accuracy(predictions, labels):
    batch_size = predictions.shape[0]
    predicted_classes = predictions.argmax(1, keepdim=True)
    correct_predictions = predicted_classes.eq(labels.view_as(predicted_classes)).sum()
    accuracy = correct_predictions.float() / batch_size
    return accuracy


# + colab={"base_uri": "https://localhost:8080/"} id="y0T-FtNN0PO3" outputId="4781c655-bd66-4326-df6d-0d9dbedbaac5"
n_epochs = 10

for epoch in range(n_epochs):

    train_loss, train_acc = train(train_dataloader, model, criterion, optimizer, device)
    valid_loss, valid_acc = evaluate(valid_dataloader, model, criterion, device)

    print(f'epoch: {epoch+1}')
    print(f'train_loss: {train_loss:.3f}, train_acc: {train_acc:.3f}')
    print(f'valid_loss: {valid_loss:.3f}, valid_acc: {valid_acc:.3f}')

# + id="OD_BKHY_42XL"


StarcoderdataJupyterScriptsDedupFiltered