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

# +
from collections import defaultdict
import math
import time
import copy
import itertools
import warnings
from datetime import datetime
warnings.filterwarnings('ignore')

import numpy as np
import torch 
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchtext import data 
import torchtext
from pathlib import Path
import pandas as pd
import spacy

from torch.utils.tensorboard import SummaryWriter
# default `log_dir` is "runs" - we'll be more specific here
writer = SummaryWriter('../runs/lstm_{}'.format(datetime.now()))
# -

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)

# # Embeddings

# test embeddings
cat_embed = nn.Embedding(5, 2)
cat_tensor = torch.LongTensor([1])
cat_embed.forward(cat_tensor)

# # Sentiment data

df = pd.read_csv("../dataset/sentiment/training.1600000.processed.noemoticon.csv", header=None, 
                 engine='python', names=['polarity', 'id', 'date', 'query', 'user', 'tweet'],
                )

df.head()

df.polarity.value_counts()

# create binary encoding 
df['sentiment'] = df.polarity.map(lambda x: 1 if x == 4 else 0)
df.head()

df.describe()

# save processed data
df.to_csv('train-processed.csv', index=None)

df.sample(10000).to_csv('train-processed-sample.csv', index=None)

# # Creating dataset

LABEL = data.LabelField()
TWEET = data.Field(tokenize='spacy', lower=True)
fields = [('polarity', None), ('id', None), ('date', None),
          ('query', None), ('user', None), ('tweet', TWEET), ('sentiment', LABEL)]
twitter_dataset = data.TabularDataset(
    path='train-processed-sample.csv',format='CSV', fields=fields, skip_header=True)

# split dataset
train, val, test = twitter_dataset.split(split_ratio=[0.8, 0.1, 0.1])
print(len(train), len(val), len(test))

# view sample
vars(train.examples[7])

# ## build a vocabulary

vocab_size = 5000
TWEET.build_vocab(train, max_size=vocab_size)
LABEL.build_vocab(train)

# how big is our vocabulary
(len(TWEET.vocab), len(LABEL.vocab))

# most common words
TWEET.vocab.freqs.most_common(10)

# create a dataloader
train_loader, val_loader, test_loader = data.BucketIterator.splits(
    (train, val, test), batch_size = 32, device = device,
    sort_key = lambda x: len(x.tweet),
    sort_within_batch = False
)


# # Create LSTM model

# +
# we use embedding and LSTM modules

class LSTM(nn.Module):
    
    def __init__(self, hidden_size, embedding_dim, vocab_size):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.encoder = nn.LSTM(input_size=embedding_dim, hidden_size=hidden_size, num_layers=1)
        self.predictor = nn.Linear(hidden_size, 2)
    
    def forward(self, seq):
        output = self.embedding(seq)
        output, (hidden, _) = self.encoder(output)
        preds = self.predictor(hidden.squeeze(0))
        return preds


# -

model = LSTM(100, 300, len(TWEET.vocab))

model


# ## Training our model

def train_model(model, optimizer, criterion, data_loader, epochs):
    model = model.to(device)
    # define training variables
    since = time.time()
    best_weights = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    
    for epoch in range(1, epochs + 1):
        print('\nEpoch {}/{}'.format(epoch, epochs))
        print('-' * 60)    
        
        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0
            
            for batch in data_loader[phase]:
                tweet, target = batch.tweet.to(device), batch.sentiment.to(device)
                # zero the parameter gradients
                optimizer.zero_grad()
                
                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    # make predictions
                    outputs = model(tweet)
                    loss = criterion(outputs, target)

                    _, preds = torch.max(outputs, 1)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * tweet.size(0)
                running_corrects += preds.eq(target.view_as(preds)).cpu().sum()
                
            batch_size = len(data_loader[phase].dataset)
            epoch_loss = running_loss / batch_size
            epoch_acc = running_corrects.double().item() / batch_size

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
            # ...log the running loss to tensorboard
            writer.add_scalar('{}/loss'.format(phase), epoch_loss, epoch)
            writer.add_scalar('{}/accuracy'.format(phase), epoch_acc, epoch)
                        
            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_weights = copy.deepcopy(model.state_dict())
            

    time_elapsed = time.time() - since
    print('-' * 60)
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:.3f} %'.format(100 * best_acc))
    print('=' * 60, '\n')
    # load best weights
    model.load_state_dict(best_weights)
    model = model.to(device)
    return model

data_loader = {'train': train_loader, 'val': val_loader}
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

m = train_model(model, optimizer, criterion, data_loader, 5)


# ## Test accuracy

def test_model(model, test_loader):
    all_targets = []
    all_preds = []
    model.eval()
    
    for batch in test_loader:
        tweets, targets = batch.tweet.to(device), batch.sentiment.to(device)
        
        outputs = model(tweets)
        _, preds = torch.max(outputs, 1)
        
        all_targets.extend(targets.cpu().numpy())
        all_preds.extend(preds.cpu().numpy())
        
    return all_targets, all_preds


y_true, y_preds = test_model(m, test_loader)

from sklearn.metrics import accuracy_score, plot_confusion_matrix, classification_report
categories = {0: 'Negative', 1: "Positive"}

score = accuracy_score(y_true, y_preds)
print("Test classification accuracy: {:.3f} %".format(score * 100))

print(classification_report(y_true, y_preds, labels=[0, 1]))


# ## classify tweets

def classify_tweet(model, tweet):
    
    processed = TWEET.process([TWEET.preprocess(tweet)])
    
    return categories[model(processed).argmax().item()]


example = test.examples[5]
vars(example)

classify_tweet(m, example.tweet)



StarcoderdataJupyterScriptsDedupFiltered