use_model.py

#!/usr/bin/env python

import itertools

import tensorflow as tf
import keras

import util

MAX_FILES_TO_USE = 100
MAX_EXAMPLES_TO_PRINT = 200


def texts_to_classes_or_subclasses(texts, lcc_class=None):
    # Note: we are not pre-loading all models in memory because they would not fit in the GPU memory
    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 = []
    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 texts_to_classes_and_subclasses(texts):
    results = texts_to_classes_or_subclasses(texts)
    for c in util.LCC_CLASSES:
        indices = []
        for i, r in enumerate(results):
            if r[0] == c:
                indices.append(i)
        if len(indices) == 0:
            continue
        texts_for_c = []
        for i in indices:
            texts_for_c.append(texts[i])
        subclasses_results = texts_to_classes_or_subclasses(texts_for_c, c)
        for i, index in enumerate(indices):
            lcc_subclass, subclass_p = subclasses_results[i]
            lcc_class, p = results[index]
            results[index] = lcc_class, p, lcc_subclass, subclass_p
    return results

def main():
    print(f"\nSelect {MAX_EXAMPLES_TO_PRINT} examples without a call number")
    docs_to_use = list(itertools.islice(util.get_documents(MAX_FILES_TO_USE, with_call_numbers=False), MAX_EXAMPLES_TO_PRINT))
    print("\nCompute the classes using the models")
    texts = []
    for doc in docs_to_use:
        texts.append(util.doc_to_text(doc))
    results = texts_to_classes_and_subclasses(texts)
    print(f"\nDisplay {MAX_EXAMPLES_TO_PRINT} examples")
    for i, doc in enumerate(docs_to_use):
        lcc_class, p, lcc_subclass, subclass_p = results[i]
        topics = ''
        if 'topic' in doc:
            topics = '- topics: ' + ' '.join(doc['topic'])
        print(f"{doc['id']} - title: {doc['title']} {topics} -> {lcc_class} ({100*p:4.1f}%) {lcc_subclass} ({100*subclass_p:4.1f}%)")

    print("\nCustom text")
    title = input("Enter a title: ")
    results = texts_to_classes_and_subclasses([title])
    lcc_class, p, lcc_subclass, subclass_p = results[0]
    print(f"-> {lcc_class} ({100*p:4.1f}%) {lcc_subclass} ({100*subclass_p:4.1f}%)")

main()