from flair.models import SequenceTagger
from nlu_utils import predict_single, predict_multiple, predict_and_annotate

# Exploratory tests
frame_model = SequenceTagger.load('frame-model/best-model.pt')
tests = [
    'chciałbym zamówić pizzę',
    'na godzinę 12',
    'prosiłbym o pizzę z pieczarkami',
    'to wszystko, jaka cena?',
    'ile kosztuje pizza',
    'do widzenia',
    'tak',
    'nie dziękuję',
    'dodatkowy ser',
    'pizzę barcelona bez cebuli',
]

# print("=== Exploratory tests - frame model ===")
for test in tests:
    print(f"Sentence: {test}")
    print(f"Single prediction: {predict_single(frame_model, test.split(), 'frame')}")
    print(f"Multiple predictions: {predict_multiple(frame_model, test.split(), 'frame')}")
    print(f"Annotated sentence: {predict_and_annotate(frame_model, test.split(), 'frame')}")

print("=== Exploratory tests - slot model ===")
slot_model = SequenceTagger.load('slot-model/final-model.pt')
for test in tests:
    print(f"Sentence: {test}")
    print(f"Prediction: {predict_and_annotate(slot_model, test.split(), 'slot')}")