generate.py

#!/usr/bin/env python

import time

import tensorflow as tf
import keras

import util

MAX_FILES_TO_USE = 100
BATCH_SIZE = 1024
P_THRESHOLD = 0.2


def read_documents():
    print("\nRead documents, prepare input text\n")
    docs = util.get_documents(MAX_FILES_TO_USE, with_call_numbers=False)
    rdocs = []
    for doc in docs:
        rdocs.append({"id": doc['id'], "txt": util.doc_to_text(doc)})
    return rdocs


def chunks(lst, n):
    # could be replaced by itertools.batched with python 3.12
    for i in range(0, len(lst), n):
        yield lst[i:i + n]


def texts_to_classes_or_subclasses(texts, model, lcc_class=None):
    results = []
    tensor = tf.convert_to_tensor(texts)
    logits = model(tensor)
    for logit in logits:
        class_idx = tf.math.argmax(logit).numpy()
        p = tf.nn.softmax(logit).numpy()[class_idx]
        if lcc_class is None:
            inferred_lcc_class = util.int_to_lcc_class(class_idx)
            results.append((inferred_lcc_class, p))
        else:
            lcc_subclass = util.int_to_lcc_subclass(lcc_class, class_idx)
            results.append((lcc_subclass, p))
    return results


def batch_classes_or_subclasses(texts, lcc_class=None):
    if lcc_class is None:
        model = keras.models.load_model('keras_models/lcc_classes.keras')
    else:
        model = keras.models.load_model(f'keras_models/lcc_class_{lcc_class}.keras')
    results = []
    n = 0
    for batch in chunks(texts, BATCH_SIZE):
        n = n + 1
        print(f"Batch number {n} / {int(len(texts)/BATCH_SIZE)+1}", end='\r')
        results.extend(texts_to_classes_or_subclasses(batch, model, lcc_class))
    keras.backend.clear_session(free_memory=True)
    return results


def compute_classes(rdocs):
    print("\nCompute classes\n")
    texts = []
    for rdoc in rdocs:
        texts.append(rdoc['txt'])
    return batch_classes_or_subclasses(texts)


def sort_documents_by_class(rdocs, lcc_classes_results):
    print("\nSort documents by class\n")
    rdocs_by_class = {}
    for c in util.LCC_CLASSES:
        rdocs_by_class[c] = []
    for i, rdoc in enumerate(rdocs):
        (c, p) = lcc_classes_results[i]
        if p > P_THRESHOLD and c in util.LCC_CLASSES:
            rdocs_by_class[c].append(rdoc)
    return rdocs_by_class


def compute_subclasses(rdocs_by_class):
    print("\nCompute subclasses\n")
    results = []
    for c in util.LCC_CLASSES:
        print(f"\nLCC class: {c}")
        texts = []
        rdocs = rdocs_by_class[c]
        for rdoc in rdocs:
            texts.append(rdoc['txt'])
        lcc_subclasses_results = batch_classes_or_subclasses(texts, c)
        for i, rdoc in enumerate(rdocs):
            (sub, p) = lcc_subclasses_results[i]
            if p > P_THRESHOLD:
                rdoc['cn'] = c + sub
                del rdoc['txt']
                results.append(rdoc)
    return results


def write_results_to_file(rdocs):
    filename = "call_numbers.csv"
    print(f"\nWrite the results into {filename}\n")
    with open(filename, 'w', encoding='utf-8') as f:
        for rdoc in rdocs:
            f.write(f"{rdoc['id']},{rdoc['cn']}\n")


def main():
    start = time.perf_counter()

    rdocs = read_documents()

    lcc_classes_results = compute_classes(rdocs)

    rdocs_by_class = sort_documents_by_class(rdocs, lcc_classes_results)

    rdocs = compute_subclasses(rdocs_by_class)

    write_results_to_file(rdocs)

    end = time.perf_counter()
    print(f"\nDone. Execution time: {int(end - start)} s")


main()