Bi_LSTM.py

import random
from collections import Counter

import numpy as np
import spacy
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.preprocessing import LabelEncoder
from torch.utils.data import DataLoader, Dataset

from base_model import BaseModel
from data_processing import CLEANUP_REGEXES, PreProcessData
from utils import group_predictions_by_position

print("PyTorch Version: ", torch.__version__)

# Comment this line if you still want to run the code without a GPU
assert torch.cuda.is_available(), "GPU is not enabled"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)


def tokenize_with_spacy(text, nlp):
    doc = nlp(text)
    tokens = []
    for token in doc:
        cleaned_token = token.text
        for regex, replacement in CLEANUP_REGEXES:
            cleaned_token = regex.sub(replacement, cleaned_token)
        cleaned_token = cleaned_token.strip()

        if cleaned_token:
            tokens.append(
                (cleaned_token, token.idx, token.idx + len(cleaned_token), token.pos_, token.tag_, token.vector)
            )
    return tokens


def create_iob_labels(tokens, annotations):
    labels = ['O'] * len(tokens)
    for ann in annotations:
        for i, (_, start, end, _, _, _) in enumerate(tokens):
            if start == ann.start:
                labels[i] = f"B-{ann.labels[0]}"
            elif ann.start < start < ann.end:
                labels[i] = f"I-{ann.labels[0]}"
    return labels


class BiLSTM_NegUnc_Dataset(Dataset):
    def __init__(self, texts, results, spacy_model, vocab=None):

        self.texts = texts
        self.results = results
        self.nlp = spacy_model

        labels = ["O", "B-NEG", "I-NEG", "B-UNC", "I-UNC", "B-NSCO", "I-NSCO", "B-USCO", "I-USCO"]

        self.label_encoder = LabelEncoder()
        self.label_encoder.fit(labels)
        print("Label encoder classes:", self.label_encoder.classes_)

        self.sentences = []
        self.labels = []

        for text, annotations in zip(self.texts, self.results):
            tokens = tokenize_with_spacy(text, self.nlp)
            iob = create_iob_labels(tokens, annotations)
            self.sentences.append(tokens)
            self.labels.append(iob)

        self.vocab = vocab if vocab else self.build_vocab()
        self.max_len = self._compute_max_len()

        print(f"Vocab size: {len(self.vocab)}")
        print(f"Max sequence length: {self.max_len}")

    def _compute_max_len(self):

        max_len = 0
        for tokens in self.sentences:
            max_len = max(max_len, len(tokens))
        return max_len

    def build_vocab(self):

        vocab = {'<PAD>': 0, '<UNK>': 1}
        idx = 2

        for tokens in self.sentences:
            for token in tokens:
                if token[0] not in vocab:
                    vocab[token[0]] = idx
                    idx += 1
        return vocab

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):

        tokens = self.sentences[idx]
        labels = self.labels[idx]

        token_texts = [token[0] for token in tokens]

        assert len(token_texts) == len(labels), f"Token: {len(token_texts)}, Label: {len(labels)}"

        encoded_labels = self.label_encoder.transform(labels)

        input_ids = [self.vocab.get(token, self.vocab.get('<UNK>', 1)) for token in token_texts]

        input_ids = input_ids[:self.max_len]
        encoded_labels = encoded_labels[:self.max_len]

        actual_len = len(input_ids)

        padding_len = self.max_len - actual_len
        if padding_len > 0:
            input_ids.extend([self.vocab['<PAD>']] * padding_len)
            encoded_labels = list(encoded_labels)
            encoded_labels.extend([0] * padding_len)

        input_ids = torch.tensor(input_ids, dtype=torch.long)
        encoded_labels = torch.tensor(encoded_labels, dtype=torch.long)

        return {
            "input_ids": input_ids,
            "labels": encoded_labels
        }


class ModelRNN(nn.Module):
    def __init__(self, input_dim, embedding_dim, hidden_dim, output_dim, n_layers, drop_prob=0.3):
        super(ModelRNN, self).__init__()
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers

        self.embedding = nn.Embedding(input_dim, embedding_dim)

        self.rnn = nn.LSTM(embedding_dim, hidden_dim, n_layers, batch_first=True, dropout=drop_prob, bidirectional=True)

        self.fc = nn.Linear(hidden_dim * 2, output_dim)

    def forward(self, x, h):
        emb = self.embedding(x)
        out, h = self.rnn(emb, h)
        out = self.fc(out)
        return out, h

    def init_hidden(self, batch_size):
        weight = next(self.parameters()).data

        if self.rnn.bidirectional:
            hidden = (weight.new_zeros(self.n_layers * 2, batch_size, self.hidden_dim),
                      weight.new_zeros(self.n_layers * 2, batch_size, self.hidden_dim))
        else:
            hidden = (weight.new_zeros(self.n_layers, batch_size, self.hidden_dim),
                      weight.new_zeros(self.n_layers, batch_size, self.hidden_dim))
        return hidden


class BiLSTM(BaseModel):
    def __init__(self, data: PreProcessData, embedding_dim, hidden_dim, n_layers):
        super().__init__(data)

        spacy.cli.download("es_core_news_md")
        self.nlp = spacy.load("es_core_news_md")

        print("Creating dataset...")
        self.dataset = BiLSTM_NegUnc_Dataset(self.data.text, self.data.results, spacy_model=self.nlp)
        self.train_loader = DataLoader(self.dataset, batch_size=32, shuffle=True)

        self.label_encoder = self.dataset.label_encoder
        self.vocab = self.dataset.vocab
        self.max_len = self.dataset.max_len
        self.weight_tensor = None

        self.input_dim = len(self.dataset.vocab)
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.output_dim = len(self.dataset.label_encoder.classes_)
        self.n_layers = n_layers

        self.model = ModelRNN(self.input_dim, embedding_dim, hidden_dim, self.output_dim, n_layers, 0.3).to(device)

    def compute_label_weights(self):
        dataset_labels = [label for sublist in self.dataset.labels for label in sublist]

        label_counts = Counter(dataset_labels)
        total_labels = len(dataset_labels)
        label_frequencies = {label: count / total_labels for label, count in label_counts.items()}
        weights = {label: 1.0 / freq for label, freq in label_frequencies.items()}

        total_weight = sum(weights.values())
        normalized_weights = {label: weight / total_weight for label, weight in weights.items()}

        labels = self.dataset.label_encoder.classes_
        self.weight_tensor = torch.tensor([normalized_weights.get(label, 1.0) for label in labels], dtype=torch.float32)

        print("Label weights:", self.weight_tensor)

    def train_model(self, lr=0.001, epochs=10, save_model=False):

        criterion = nn.CrossEntropyLoss(weight=self.weight_tensor.to(device) if self.weight_tensor else None)
        optimizer = optim.Adam(self.model.parameters(), lr=lr)

        for epoch in range(epochs):
            self.model.train()
            running_loss = 0

            for batch_idx, batch in enumerate(self.train_loader):
                input_ids = batch['input_ids'].to(device)
                labels = batch['labels'].to(device)

                optimizer.zero_grad()
                h = self.model.init_hidden(input_ids.size(0))

                outputs, h = self.model(input_ids, h)

                outputs = outputs.view(-1, self.output_dim)
                labels = labels.view(-1)

                loss = criterion(outputs, labels)

                _, predictions = torch.max(outputs, dim=1)
                loss.backward()
                optimizer.step()
                running_loss += loss.item()

            print(f'Epoch {epoch + 1}/{epochs}, Loss: {running_loss / len(self.train_loader)}')

        if save_model:
            self.save_model()

    def predict(self, text):
        self.model.eval()
        tokens = tokenize_with_spacy(text, self.nlp)
        input_ids = torch.tensor(
            [self.vocab.get(token[0], self.vocab.get('<UNK>', 1)) for token in tokens][:self.max_len], dtype=torch.long)

        input_ids = input_ids.unsqueeze(0).to(device)

        h = self.model.init_hidden(input_ids.size(0))

        with torch.no_grad():
            outputs, h = self.model(input_ids, h)

        probabilities = torch.softmax(outputs, dim=2)

        predicted_labels = torch.argmax(probabilities, dim=2).squeeze().cpu().numpy()

        decoded_labels = self.label_encoder.inverse_transform(predicted_labels)

        results = []
        for (tok, start, end, _, _, _), label in zip(tokens, decoded_labels):
            if label != 'O' and tok and tok.strip():
                tag = label.split('-')[-1]
                results.append({
                    'text': tok,
                    'label': tag,
                    'start': start,
                    'end': end
                })

        preds = group_predictions_by_position(results)
        return preds

    def save_model(self, path="model_weights.pth"):
        torch.save(self.model.state_dict(), path)
        print(f"Model saved to {path}")

    def load_model(self, path="model_weights.pth"):
        self.model.load_state_dict(torch.load(path, map_location=device))
        self.model.to(device)
        self.model.eval()
        print(f"Model loaded from {path}")