hing_bert_module/utils.py

import os
from typing import Sequence
from .classifier import TokenPrediction

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
OUTPUT_PATH = os.path.join(BASE_DIR, 'output2.txt')
LOG_INITIALIZED = False


def _print_predictions(predictions: Sequence[TokenPrediction]):
    print("Token\tLabel\tConfidence")
    for pred in predictions:
        print(f"{pred.token}\t{pred.label}\t{pred.confidence:.4f}")


def _init_output_log():
    global LOG_INITIALIZED
    if LOG_INITIALIZED:
        return
    with open(OUTPUT_PATH, 'w', encoding='utf-8') as f:
        f.write('HingBERT-LID session log\n')
    LOG_INITIALIZED = True


def _write_predictions(predictions: Sequence[TokenPrediction], source_text: str):
    """Write predictions to output file."""
    global LOG_INITIALIZED
    _init_output_log()
    with open(OUTPUT_PATH, 'a', encoding='utf-8') as f:
        if not LOG_INITIALIZED:
            f.write('\n' + '='*80 + '\n')
            LOG_INITIALIZED = True
        f.write(f'\nSource: {source_text}\n')
        for pred in predictions:
            f.write(f"{pred.token}\t{pred.label}\t{pred.confidence:.4f}\n")