chat-restaruacja/nlu_tests.py

30 lines
1.1 KiB
Python

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')}")