run.py

#############################################################################
### Търсене и извличане на информация. Приложение на дълбоко машинно обучение
### Стоян Михов
### Зимен семестър 2025/2026
#############################################################################
###
### Машинен превод чрез генеративен езиков модел
###
#############################################################################
from utils.train import get_lr
from bpe.tokenizer import encode_word
from bpe.loader import get_bpe, build_bpe_lookup
from parameters.model import ModelParams
from utils.meta import progress_bar
from utils.file import save_to_file, load_from_file
from utils.corpus import prepare_data, read_corpus_from_disk, get_train_corpus
import sys
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import math
import time
from nltk.translate.bleu_score import corpus_bleu
from model import DecoderTransformer
from parameters.nlp import (
    UNK_TOKEN_IDX,
    TRANS_TOKEN_IDX,
    PAD_TOKEN_IDX,
    START_TOKEN_IDX,
    END_TOKEN_IDX,
    SPECIAL_TOKENS,
)
from parameters.hyper import (
    MAX_EPOCHS,
    LEARNING_RATE,
    DEVICE,
    BATCH_SIZE,
    CLIP_GRAD,
    LOG_EVERY,
    TEST_EVERY,
)
from parameters.corpus import (
    WORDS_FILE_NAME,
    CORPUS_FILE_NAME,
    MODEL_FILE_NAME,
)


def perplexity(nmt, test, BATCH_SIZE):
    test_size = len(test)
    H = 0.0
    c = 0
    for b in range(0, test_size, BATCH_SIZE):
        batch = test[b : min(b + BATCH_SIZE, test_size)]

        padded = prepare_padded_batch(batch, DEVICE)

        c += (padded[:, 1:] != PAD_TOKEN_IDX).sum().item()
        with torch.no_grad():
            # H += L * nmt(batch)
            logits = nmt(padded[:, :-1])

            H += F.cross_entropy(
                logits.transpose(
                    -1, -2
                ),  # \in \mathbb{R}^{B \times N-1 \times K} but cross entropy wants the classes second
                padded[:, 1:],
                ignore_index=PAD_TOKEN_IDX,
                reduction="sum",
            ).item()
    return math.exp(H / c)


def prepare_padded_batch(tokens: list[list[int]], device):
    m = max(len(tl) for tl in tokens)
    sents_padded = [s + (m - len(s)) * [PAD_TOKEN_IDX] for s in tokens]
    return torch.tensor(sents_padded, dtype=torch.long, device=device)


class RuntimeUtils:
    @staticmethod
    def prepare():
        (train_corpus, dev_corpus), vocab = prepare_data()
        save_to_file((train_corpus, dev_corpus), CORPUS_FILE_NAME)
        save_to_file(vocab, WORDS_FILE_NAME)
        print("Data prepared.")

    @staticmethod
    def test_prepare():
        (train_corpus, dev_corpus) = load_from_file(CORPUS_FILE_NAME)
        vocab = load_from_file(WORDS_FILE_NAME)

        print("Loaded prepared data:")
        print(train_corpus[:5])
        print(dev_corpus[:5])
        print(list(vocab.items())[:100])

    @staticmethod
    def train(argv_1: str):
        (train_corpus, dev_corpus) = load_from_file(CORPUS_FILE_NAME)
        vocab = load_from_file(WORDS_FILE_NAME)

        model_params = ModelParams(N=len(vocab))
        learning_rate = LEARNING_RATE
        nmt = DecoderTransformer(model_params).to(DEVICE)
        optimizer = torch.optim.AdamW(nmt.parameters(), lr=learning_rate)

        if argv_1 == "extratrain":
            nmt.load(MODEL_FILE_NAME)
            (iter, best_perplexity, learning_rate, osd) = torch.load(
                f"{MODEL_FILE_NAME}.optim"
            )
            optimizer.load_state_dict(osd)
            for param_group in optimizer.param_groups:
                param_group["lr"] = learning_rate
        else:
            best_perplexity = math.inf
            iter = 0

        idx = np.arange(len(train_corpus), dtype="int32")
        nmt.train()
        beginTime = time.time()
        for epoch in range(MAX_EPOCHS):
            np.random.shuffle(idx)
            words = 0
            train_time = time.time()
            for b in range(0, len(idx), BATCH_SIZE):
                ######################################################################################################
                ### Може да се наложи да се променя скоростта на спускане learning_rate в зависимост от итерацията
                ######################################################################################################
                iter += 1
                lr = get_lr(iter)
                for param_group in optimizer.param_groups:
                    param_group["lr"] = lr
                batch = [
                    train_corpus[i] for i in idx[b : min(b + BATCH_SIZE, len(idx))]
                ]

                words += sum(len(s) - 1 for s in batch)
                padded = prepare_padded_batch(batch, DEVICE)
                logits = nmt(padded[:, :-1])
                H = F.cross_entropy(
                    logits.transpose(
                        -1, -2
                    ),  # \in \mathbb{R}^{B \times N-1 \times K} but cross entropy wants the classes second
                    padded[:, 1:],
                    ignore_index=PAD_TOKEN_IDX,
                )
                optimizer.zero_grad()
                H.backward()
                nn.utils.clip_grad_norm_(nmt.parameters(), CLIP_GRAD)
                optimizer.step()

                if iter % LOG_EVERY == 0:
                    print(
                        "Iteration:",
                        iter,
                        "Epoch:",
                        epoch + 1,
                        "/",
                        MAX_EPOCHS,
                        ", Batch:",
                        b // BATCH_SIZE + 1,
                        "/",
                        len(idx) // BATCH_SIZE + 1,
                        ", loss: ",
                        H.item(),
                        "words/sec:",
                        words / (time.time() - train_time),
                        "time elapsed:",
                        (time.time() - beginTime),
                    )
                    train_time = time.time()
                    words = 0

                if iter % TEST_EVERY == 0:
                    nmt.eval()
                    currentPerplexity = perplexity(nmt, dev_corpus, BATCH_SIZE)
                    nmt.train()
                    print("Current model perplexity: ", currentPerplexity)

                    if currentPerplexity < best_perplexity:
                        best_perplexity = currentPerplexity
                        print("Saving new best model.")
                        nmt.save(MODEL_FILE_NAME)
                        torch.save(
                            (
                                iter,
                                best_perplexity,
                                learning_rate,
                                optimizer.state_dict(),
                            ),
                            f"{MODEL_FILE_NAME}.optim",
                        )

        print("Reached maximum number of epochs!")
        nmt.eval()
        currentPerplexity = perplexity(nmt, dev_corpus, BATCH_SIZE)
        print("Last model perplexity: ", currentPerplexity)

        if currentPerplexity < best_perplexity:
            best_perplexity = currentPerplexity
            print("Saving last model.")
            nmt.save(MODEL_FILE_NAME)
            torch.save(
                (iter, best_perplexity, learning_rate, optimizer.state_dict()),
                f"{MODEL_FILE_NAME}.optim",
            )

    @staticmethod
    def perplexity(argv_2: str, argv_3: str):
        vocab = load_from_file(WORDS_FILE_NAME)
        source_corpus, target_corpus = get_train_corpus()
        (bpe_vocab, bpe_merges) = get_bpe(source_corpus, target_corpus)
        bpe_lookup = build_bpe_lookup(bpe_vocab)

        model_params = ModelParams(N=len(vocab))
        nmt = DecoderTransformer(model_params).to(DEVICE)
        nmt.load(MODEL_FILE_NAME)

        source_test = read_corpus_from_disk(argv_2)
        target_test = read_corpus_from_disk(argv_3)

        test = []
        test_corpus = zip(source_test, target_test)
        for source_s, target_s in test_corpus:
            sent = [START_TOKEN_IDX]

            for w in source_s:
                if w in bpe_lookup:
                    ew = bpe_lookup[w]
                else:
                    ew = encode_word(w, bpe_merges)
                ew = ew.split(" ")
                for t in ew:
                    sent.append(vocab.get(t, UNK_TOKEN_IDX))

            sent.append(TRANS_TOKEN_IDX)

            for w in target_s:
                if w in bpe_lookup:
                    ew = bpe_lookup[w]
                else:
                    ew = encode_word(w, bpe_merges)
                ew = ew.split(" ")
                for t in ew:
                    sent.append(vocab.get(t, UNK_TOKEN_IDX))

            sent.append(END_TOKEN_IDX)

            test.append(sent)

        nmt.eval()
        print("Model perplexity: ", perplexity(nmt, test, BATCH_SIZE))

    @staticmethod
    @progress_bar
    def translate(argv_2: str, argv_3: str, **kwargs):
        vocab = load_from_file(WORDS_FILE_NAME)
        source_corpus, target_corpus = get_train_corpus()
        (bpe_vocab, bpe_merges) = get_bpe(source_corpus, target_corpus)
        bpe_lookup = build_bpe_lookup(bpe_vocab)
        words = list(vocab)

        source_test = read_corpus_from_disk(argv_2)

        test = []
        for s in source_test:
            sent = [START_TOKEN_IDX]
            for w in s:
                if w in bpe_lookup:
                    ew = bpe_lookup[w]
                else:
                    ew = encode_word(w, bpe_merges)
                ew = ew.split(" ")

                for t in ew:
                    sent.append(
                        vocab.get(
                            t,
                            UNK_TOKEN_IDX,
                        )
                    )
            sent.append(TRANS_TOKEN_IDX)
            test.append(sent)

        model_params = ModelParams(N=len(vocab))
        nmt = DecoderTransformer(model_params).to(DEVICE)
        nmt.load(MODEL_FILE_NAME)

        nmt.eval()
        with open(argv_3, "w") as file:
            kwargs["pb_init"](len(test))
            for s in test:
                r = nmt.generate(s)
                st = r.index(TRANS_TOKEN_IDX)
                result = [words[i] for i in r[st + 1 : -1]]
                file.write("".join(result).replace("</w>", " ") + "\n")
                kwargs["pb_tick"]()

    @staticmethod
    def generate(argv_2: str):
        vocab = load_from_file(WORDS_FILE_NAME)
        source_corpus, target_corpus = get_train_corpus()
        (bpe_vocab, bpe_merges) = get_bpe(source_corpus, target_corpus)
        bpe_lookup = build_bpe_lookup(bpe_vocab)
        words = list(vocab)

        test = []
        for w in argv_2.split():
            if w in SPECIAL_TOKENS:
                ew = w
            elif w in bpe_lookup:
                ew = bpe_lookup[w]
            else:
                ew = encode_word(w, bpe_merges)
            ew = ew.split(" ")

            for t in ew:
                test.append(
                    vocab.get(
                        t,
                        UNK_TOKEN_IDX,
                    )
                )

        model_params = ModelParams(N=len(vocab))
        nmt = DecoderTransformer(model_params).to(DEVICE)
        nmt.load(MODEL_FILE_NAME)

        nmt.eval()
        r = nmt.generate(test)
        result = [words[i] for i in r]
        print("".join(result).replace("</w>", " ") + "\n")

    @staticmethod
    def bleu(argv_2: str, argv_3: str):
        ref = [[s] for s in read_corpus_from_disk(argv_2)]
        hyp = read_corpus_from_disk(argv_3)

        bleu_score = corpus_bleu(ref, hyp)
        print("Corpus BLEU: ", (bleu_score * 100))


if len(sys.argv) > 1 and sys.argv[1] == "prepare":
    RuntimeUtils.prepare()
elif len(sys.argv) > 1 and sys.argv[1] == "test_prepare":
    RuntimeUtils.test_prepare()
elif len(sys.argv) > 1 and sys.argv[1] in ("train", "extratrain"):
    RuntimeUtils.train(sys.argv[1])
elif len(sys.argv) > 3 and sys.argv[1] == "perplexity":
    RuntimeUtils.perplexity(sys.argv[2], sys.argv[3])
elif len(sys.argv) > 3 and sys.argv[1] == "translate":
    RuntimeUtils.translate(sys.argv[2], sys.argv[3])
elif len(sys.argv) > 2 and sys.argv[1] == "generate":
    RuntimeUtils.generate(sys.argv[2])
elif len(sys.argv) > 3 and sys.argv[1] == "bleu":
    RuntimeUtils.bleu(sys.argv[2], sys.argv[3])