Embeddings Py 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)
print("Vocab size = ", vocab_size)

class EmbedClassifier(torch.nn.Module):
    def __init__(self, vocab_size, embed_dim, num_class):
        super().__init__()
        self.embedding = torch.nn.Embedding(vocab_size, embed_dim)
        self.fc = torch.nn.Linear(embed_dim, num_class)

    def forward(self, x):
        x = self.embedding(x)
        x = torch.mean(x,dim=1)
        return self.fc(x)

def padify(b):
    # b is the list of tuples of length batch_size
    #   - first element of a tuple = label,
    #   - second = feature (text sequence)
    # build vectorized sequence
    v = [encode(x[1]) for x in b]
    # first, compute max length of a sequence in this minibatch
    l = max(map(len,v))
    return ( # tuple of two tensors - labels and features
        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])
    )

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, collate_fn=padify, shuffle=True)

net = EmbedClassifier(vocab_size,32,len(classes)).to(device)
train_epoch(net,train_loader, lr=1, epoch_size=25000)

class EmbedClassifier(torch.nn.Module):
    def __init__(self, vocab_size, embed_dim, num_class):
        super().__init__()
        self.embedding = torch.nn.EmbeddingBag(vocab_size, embed_dim)
        self.fc = torch.nn.Linear(embed_dim, num_class)

    def forward(self, text, off):
        x = self.embedding(text, off)
        return self.fc(x)

def offsetify(b):
    # first, compute data tensor from all sequences
    x = [torch.tensor(encode(t[1])) for t in b]
    # now, compute the offsets by accumulating the tensor of sequence lengths
    o = [0] + [len(t) for t in x]
    o = torch.tensor(o[:-1]).cumsum(dim=0)
    return (
        torch.LongTensor([t[0]-1 for t in b]), # labels
        torch.cat(x), # text
        o
    )

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, collate_fn=offsetify, shuffle=True)

net = EmbedClassifier(vocab_size,32,len(classes)).to(device)

def train_epoch_emb(net,dataloader,lr=0.01,optimizer=None,loss_fn = torch.nn.CrossEntropyLoss(),epoch_size=None, report_freq=200):
    optimizer = optimizer or torch.optim.Adam(net.parameters(),lr=lr)
    loss_fn = loss_fn.to(device)
    net.train()
    total_loss,acc,count,i = 0,0,0,0
    for labels,text,off in dataloader:
        optimizer.zero_grad()
        labels,text,off = labels.to(device), text.to(device), off.to(device)
        out = net(text, off)
        loss = loss_fn(out,labels) #cross_entropy(out,labels)
        loss.backward()
        optimizer.step()
        total_loss+=loss
        _,predicted = torch.max(out,1)
        acc+=(predicted==labels).sum()
        count+=len(labels)
        i+=1
        if i%report_freq==0:
            print(f"{count}: acc={acc.item()/count}")
        if epoch_size and count>epoch_size:
            break
    return total_loss.item()/count, acc.item()/count


train_epoch_emb(net,train_loader, lr=4, epoch_size=25000)

COURSE_SEMANTIC_GROUPS = [
    ("royalty", {"king", "queen"}),
    ("masculine", {"king", "man"}),
    ("feminine", {"queen", "woman"}),
    ("sports", {"athletes", "basketball", "coach", "football", "play", "runner", "swimming", "team", "tennis", "volleyball"}),
    ("business", {"bank", "business", "companies", "contracts", "funds", "market", "payment", "prices", "revenue", "sales"}),
    ("technology", {"battery", "cloud", "engineers", "model", "neural", "researchers", "robot", "satellite", "scientists", "sensor", "software", "technology"}),
    ("world", {"council", "countries", "diplomats", "election", "ministers", "paris", "regional", "summit"}),
    ("care", {"emergency", "health", "hospital", "humanitarian", "medicine", "relief", "safety"}),
]

def build_course_semantic_vectors(vocabulary):
    return [
        [1.0 if word in members else 0.0 for _name, members in COURSE_SEMANTIC_GROUPS]
        for word in vocabulary.itos
    ]

def course_cosine_similarity(left, right):
    numerator = sum(a * b for a, b in zip(left, right))
    left_norm = sum(value * value for value in left) ** 0.5
    right_norm = sum(value * value for value in right) ** 0.5
    if left_norm == 0.0 or right_norm == 0.0:
        return -1.0
    return numerator / (left_norm * right_norm)

class CourseSemanticVectors:
    def __init__(self, vocabulary):
        self.itos = vocabulary.get_itos()
        self.stoi = vocabulary.get_stoi()
        self.vectors = build_course_semantic_vectors(vocabulary)

    def get_vector(self, word):
        if word not in self.stoi:
            raise ValueError(f"{word!r} is not in the course vocabulary")
        return list(self.vectors[self.stoi[word]])

    def most_similar(self, word=None, positive=None, negative=None, topn=5):
        positive = list(positive or [])
        negative = list(negative or [])
        if word is not None:
            positive.insert(0, word)
        if not positive:
            raise ValueError("provide a word or at least one positive token")

        query = [0.0] * len(COURSE_SEMANTIC_GROUPS)
        for token in positive:
            query = [a + b for a, b in zip(query, self.get_vector(token))]
        for token in negative:
            query = [a - b for a, b in zip(query, self.get_vector(token))]

        excluded = set(positive + negative)
        ranked = [
            (token, course_cosine_similarity(query, vector))
            for token, vector in zip(self.itos, self.vectors)
            if token not in excluded and any(vector)
        ]
        ranked.sort(key=lambda item: (-item[1], item[0]))
        return ranked[:topn]

course_semantics = CourseSemanticVectors(vocab)
print("Authored semantic dimensions:", [name for name, _members in COURSE_SEMANTIC_GROUPS])

for word, similarity in course_semantics.most_similar("neural"):
    print(f"{word} -> {similarity:.3f}")

course_semantics.get_vector("play")

course_semantics.most_similar(
    positive=["king", "woman"],
    negative=["man"],
    topn=1,
)[0]

course_vector_matrix = torch.tensor(course_semantics.vectors, dtype=torch.float32)
embed_size = course_vector_matrix.shape[1]
print(f"Embedding size: {embed_size}")

net = EmbedClassifier(vocab_size, embed_size, len(classes))
with torch.no_grad():
    net.embedding.weight.copy_(course_vector_matrix)
net = net.to(device)

train_epoch_emb(net, train_loader, lr=1, epoch_size=25000)

def build_cooccurrence_embeddings(dataset, vocabulary, window_size=2, dimensions=16):
    matrix = torch.zeros((len(vocabulary), len(vocabulary)), dtype=torch.float64)
    for _label, text in dataset:
        token_ids = encode(text, voc=vocabulary)
        for center, center_id in enumerate(token_ids):
            start = max(0, center - window_size)
            stop = min(len(token_ids), center + window_size + 1)
            for context in range(start, stop):
                if context != center:
                    matrix[center_id, token_ids[context]] += 1.0

    matrix = torch.log1p(matrix)
    left, singular_values, _right = torch.linalg.svd(matrix, full_matrices=False)
    dimensions = min(dimensions, left.shape[1])
    left = left[:, :dimensions].clone()
    for column in range(dimensions):
        pivot = int(torch.argmax(torch.abs(left[:, column])).item())
        if left[pivot, column] < 0:
            left[:, column] *= -1
    vectors = left * torch.sqrt(singular_values[:dimensions])
    vectors[0] = 0
    return vectors.to(torch.float32)

course_vectors = build_cooccurrence_embeddings(train_dataset, vocab)
print("Teaching embedding shape:", tuple(course_vectors.shape))

def nearest_course_words(word, n=5):
    if word not in vocab.stoi:
        raise ValueError(f"{word!r} is not in the course vocabulary")
    index = vocab[word]
    distances = torch.linalg.vector_norm(course_vectors - course_vectors[index], dim=1)
    distances[index] = torch.inf
    nearest = torch.argsort(distances)[:n].tolist()
    return [(vocab.itos[item], float(distances[item])) for item in nearest]

nearest_course_words("technology")

def offsetify(b):
    # first, compute data tensor from all sequences
    x = [torch.tensor(encode(t[1],voc=vocab)) for t in b] # pass the instance of vocab to encode function!
    # now, compute the offsets by accumulating the tensor of sequence lengths
    o = [0] + [len(t) for t in x]
    o = torch.tensor(o[:-1]).cumsum(dim=0)
    return (
        torch.LongTensor([t[0]-1 for t in b]), # labels
        torch.cat(x), # text
        o
    )

net = EmbedClassifier(len(vocab), course_vectors.shape[1], len(classes))
with torch.no_grad():
    net.embedding.weight.copy_(course_vectors)
net = net.to(device)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, collate_fn=offsetify, shuffle=True)
train_epoch_emb(net,train_loader, lr=4, epoch_size=25000)