RNNPy Torch

import collections
import random
import re

import numpy as np
import torch

# course-edition deterministic NLP utilities v2

COURSE_NEWS_CLASSES = ["World", "Sports", "Business", "Sci/Tech"]

# This compact corpus was written for this edition. It keeps the notebooks
# deterministic, network-free, and quick enough to run on a learner's CPU.
COURSE_NEWS_TRAIN = [
    (1, "Paris hosted a regional summit where diplomats discussed water security, peaceful trade routes, and a shared emergency plan for neighboring communities."),
    (1, "The elected council approved a cross-border health agreement after delegates reviewed hospital capacity, medicine access, and transparent public reporting."),
    (1, "A coastal city welcomed observers for a public vote, while local groups published clear guidance about polling places and accessible transportation."),
    (1, "The king and queen met a woman who leads the relief agency and a man who coordinates volunteers, then the group announced a neutral humanitarian corridor."),
    (1, "Ministers signed a climate adaptation pledge that funds drought monitoring, resilient farms, and open scientific exchange across several countries."),
    (1, "Community mediators opened a week of talks focused on civilian safety, reliable food deliveries, and practical steps toward a lasting ceasefire."),
    (2, "The basketball team completed a patient comeback in the final quarter, using quick passes, strong defense, and a balanced scoring plan to win the tournament."),
    (2, "A young runner broke the course record after months of careful training, recovery sessions, nutrition planning, and support from her local athletics club."),
    (2, "The football coach praised disciplined play after the squad protected an early lead and created chances through accurate movement on both wings."),
    (2, "Fans filled the arena for a close volleyball final in which both teams served aggressively, defended long rallies, and respected every referee decision."),
    (2, "The tennis champion returned after injury and won a demanding match by varying pace, placing serves carefully, and staying calm during two tie breaks."),
    (2, "Organizers added inclusive swimming events to the city games so more athletes can compete with safe facilities, trained officials, and fair timing equipment."),
    (3, "Microsoft announced a small-business program that combines cloud credits, security workshops, and practical accounting support for new local companies."),
    (3, "Global funds moved cautiously after the central bank held interest rates steady and asked lenders to publish clearer information about consumer borrowing costs."),
    (3, "A neighborhood market expanded its delivery service, hired twelve workers, and invested revenue in reusable packaging supplied by another local business."),
    (3, "The manufacturer reported stable quarterly sales while noting that shipping delays and higher material prices could reduce margins later in the year."),
    (3, "Two payment companies agreed to share fraud signals through a privacy-preserving system designed to protect customers without exposing purchase details."),
    (3, "A cooperative offered farmers transparent contracts and faster invoices, helping members plan equipment repairs before the next growing season begins."),
    (4, "Researchers trained a compact neural network to identify damaged solar panels, then published the evaluation data and documented where the model still fails."),
    (4, "A university technology lab released an open sensor design that measures classroom air quality and stores only anonymous readings for public analysis."),
    (4, "Engineers improved a battery recycling process by recovering more useful minerals at lower temperatures, reducing energy use during the pilot study."),
    (4, "The space telescope captured a detailed spectrum from a distant planet, giving scientists new evidence about clouds and molecules in its atmosphere."),
    (4, "A software team tested an accessibility assistant with keyboard users and screen-reader experts before changing the interface and publishing the results."),
    (4, "Students built a small robot that maps indoor obstacles with inexpensive sensors, explains each route choice, and says hello when a test run begins."),
]

COURSE_NEWS_TEST = [
    (1, "Regional delegates published a joint disaster response schedule after reviewing evacuation routes and communication gaps."),
    (1, "Election monitors confirmed the final count and recommended clearer access rules for voters who need assistance."),
    (2, "The basketball captain scored late, but credited the victory to defense, passing, and careful preparation across the whole season."),
    (2, "A cycling club opened a safe youth race with trained marshals, marked turns, and free equipment checks."),
    (3, "Technology shares rose while retail funds remained cautious after companies issued mixed forecasts for the next quarter."),
    (3, "The family business secured a modest loan to replace old equipment and expand its apprenticeship program."),
    (4, "A research team released a smaller language model with documented energy measurements and a public evaluation set."),
    (4, "Scientists used open satellite data to improve flood warnings and shared the software with local emergency teams."),
]

def course_tokenize(text):
    return re.findall(r"[a-z0-9]+(?:'[a-z0-9]+)?", str(text).lower())

def course_ngrams_iterator(tokens, ngrams=1):
    tokens = list(tokens)
    for token in tokens:
        yield token
    for size in range(2, max(1, ngrams) + 1):
        for start in range(0, len(tokens) - size + 1):
            yield " ".join(tokens[start:start + size])

class CourseVocabulary:
    def __init__(self, counter, min_freq=1, max_tokens=None):
        ordered = sorted(
            (token for token, count in counter.items() if count >= min_freq),
            key=lambda token: (-counter[token], token),
        )
        if max_tokens is not None:
            ordered = ordered[:max(0, max_tokens - 1)]
        self.itos = ["<unk>"] + [token for token in ordered if token != "<unk>"]
        self.stoi = {token: index for index, token in enumerate(self.itos)}

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

    def __getitem__(self, token):
        return self.stoi.get(token, 0)

    def get_stoi(self):
        return dict(self.stoi)

    def get_itos(self):
        return list(self.itos)

def build_course_vocab(dataset, ngrams=1, min_freq=1, max_tokens=None, tokenizer=course_tokenize):
    counter = collections.Counter()
    for _label, text in dataset:
        counter.update(course_ngrams_iterator(tokenizer(text), ngrams=ngrams))
    return CourseVocabulary(counter, min_freq=min_freq, max_tokens=max_tokens)

def load_course_fixture(ngrams=1, min_freq=1, vocab_size=None, lines_cnt=None):
    global vocab, tokenizer
    tokenizer = course_tokenize
    train_dataset = list(COURSE_NEWS_TRAIN)
    test_dataset = list(COURSE_NEWS_TEST)
    vocabulary_rows = train_dataset if lines_cnt is None else train_dataset[:lines_cnt]
    vocab = build_course_vocab(
        vocabulary_rows,
        ngrams=ngrams,
        min_freq=min_freq,
        max_tokens=vocab_size,
        tokenizer=tokenizer,
    )
    return train_dataset, test_dataset, list(COURSE_NEWS_CLASSES), vocab

def encode(text, voc=None, unk=0, tokenizer=course_tokenize):
    selected_vocab = vocab if voc is None else voc
    stoi = selected_vocab.get_stoi()
    return [stoi.get(token, unk) for token in tokenizer(text)]

def train_epoch(net, dataloader, lr=0.01, optimizer=None, loss_fn=None, epoch_size=None, report_freq=200):
    optimizer = optimizer or torch.optim.Adam(net.parameters(), lr=lr)
    loss_fn = (loss_fn or torch.nn.CrossEntropyLoss()).to(device)
    net.train()
    total_loss, accuracy, count, batch_index = 0.0, 0, 0, 0
    for labels, features in dataloader:
        optimizer.zero_grad()
        features, labels = features.to(device), labels.to(device)
        output = net(features)
        loss = loss_fn(output, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        accuracy += (output.argmax(1) == labels).sum().item()
        count += len(labels)
        batch_index += 1
        if batch_index % report_freq == 0:
            print(f"{count}: acc={accuracy / count:.4f}")
        if epoch_size and count > epoch_size:
            break
    return total_loss / max(count, 1), accuracy / max(count, 1)

def padify(batch, voc=None, tokenizer=course_tokenize):
    vectors = [encode(item[1], voc=voc, tokenizer=tokenizer) for item in batch]
    max_length = max(map(len, vectors))
    return (
        torch.LongTensor([item[0] - 1 for item in batch]),
        torch.stack([
            torch.nn.functional.pad(torch.tensor(vector), (0, max_length - len(vector)), value=0)
            for vector in vectors
        ]),
    )

def offsetify(batch, voc=None):
    vectors = [torch.tensor(encode(item[1], voc=voc)) for item in batch]
    offsets = torch.tensor(([0] + [len(vector) for vector in vectors])[:-1]).cumsum(dim=0)
    return torch.LongTensor([item[0] - 1 for item in batch]), torch.cat(vectors), offsets

def train_epoch_emb(net, dataloader, lr=0.01, optimizer=None, loss_fn=None, epoch_size=None, report_freq=200, use_pack_sequence=False):
    optimizer = optimizer or torch.optim.Adam(net.parameters(), lr=lr)
    loss_fn = (loss_fn or torch.nn.CrossEntropyLoss()).to(device)
    net.train()
    total_loss, accuracy, count, batch_index = 0.0, 0, 0, 0
    for labels, text, offsets in dataloader:
        optimizer.zero_grad()
        labels, text = labels.to(device), text.to(device)
        offsets = offsets.to("cpu" if use_pack_sequence else device)
        output = net(text, offsets)
        loss = loss_fn(output, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        accuracy += (output.argmax(1) == labels).sum().item()
        count += len(labels)
        batch_index += 1
        if batch_index % report_freq == 0:
            print(f"{count}: acc={accuracy / count:.4f}")
        if epoch_size and count > epoch_size:
            break
    return total_loss / max(count, 1), accuracy / max(count, 1)

random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = course_tokenize
vocab = None

train_dataset, test_dataset, classes, vocab = load_course_fixture()
vocab_size = len(vocab)

class RNNClassifier(torch.nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_class):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.embedding = torch.nn.Embedding(vocab_size, embed_dim)
        self.rnn = torch.nn.RNN(embed_dim,hidden_dim,batch_first=True)
        self.fc = torch.nn.Linear(hidden_dim, num_class)

    def forward(self, x):
        batch_size = x.size(0)
        x = self.embedding(x)
        x,h = self.rnn(x)
        return self.fc(x.mean(dim=1))

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, collate_fn=padify, shuffle=True)
net = RNNClassifier(vocab_size,64,32,len(classes)).to(device)
train_epoch(net,train_loader, lr=0.001)

class LSTMClassifier(torch.nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_class):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.embedding = torch.nn.Embedding(vocab_size, embed_dim)
        self.embedding.weight.data = torch.randn_like(self.embedding.weight.data)-0.5
        self.rnn = torch.nn.LSTM(embed_dim,hidden_dim,batch_first=True)
        self.fc = torch.nn.Linear(hidden_dim, num_class)

    def forward(self, x):
        batch_size = x.size(0)
        x = self.embedding(x)
        x,(h,c) = self.rnn(x)
        return self.fc(h[-1])

net = LSTMClassifier(vocab_size,64,32,len(classes)).to(device)
train_epoch(net,train_loader, lr=0.001)

def pad_length(b):
    # build vectorized sequence
    v = [encode(x[1]) for x in b]
    # compute max length of a sequence in this minibatch and length sequence itself
    len_seq = list(map(len,v))
    l = max(len_seq)
    return ( # tuple of three tensors - labels, padded features, length sequence
        torch.LongTensor([t[0]-1 for t in b]),
        torch.stack([torch.nn.functional.pad(torch.tensor(t),(0,l-len(t)),mode='constant',value=0) for t in v]),
        torch.tensor(len_seq)
    )

train_loader_len = torch.utils.data.DataLoader(train_dataset, batch_size=16, collate_fn=pad_length, shuffle=True)

class LSTMPackClassifier(torch.nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_class):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.embedding = torch.nn.Embedding(vocab_size, embed_dim)
        self.embedding.weight.data = torch.randn_like(self.embedding.weight.data)-0.5
        self.rnn = torch.nn.LSTM(embed_dim,hidden_dim,batch_first=True)
        self.fc = torch.nn.Linear(hidden_dim, num_class)

    def forward(self, x, lengths):
        batch_size = x.size(0)
        x = self.embedding(x)
        pad_x = torch.nn.utils.rnn.pack_padded_sequence(x,lengths,batch_first=True,enforce_sorted=False)
        pad_x,(h,c) = self.rnn(pad_x)
        x, _ = torch.nn.utils.rnn.pad_packed_sequence(pad_x,batch_first=True)
        return self.fc(h[-1])

net = LSTMPackClassifier(vocab_size,64,32,len(classes)).to(device)
train_epoch_emb(net,train_loader_len, lr=0.001,use_pack_sequence=True)