File size: 1,129 Bytes
8a02978 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
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")
|