sklep-internetowy-systemy-dialogowe/chatbot/modules/nlu.py

from flair.models import SequenceTagger
import sys

sys.path.append("..")
from models.nlu_train2 import predict_frame, predict_slot
import logging

logging.getLogger('flair').setLevel(logging.CRITICAL)


class Slot:
    def __init__(self, name, value=None):
        self.name = name
        self.value = value

    def __str__(self) -> str:
        return f"Name: {self.name}, Value: {self.value}"


class Act:
    def __init__(self, intent: str, slots: list[Slot] = []):
        self.slots = slots
        self.intent = intent

    def __str__(self):
        msg = f"Act: {self.intent}, Slots: ["
        for slot in self.slots:
            msg += f"({slot}), "
        msg += "]"
        return msg


class NLU:
    def __init__(self):
        self.frame_model = SequenceTagger.load('models/frame-model/final-model.pt')
        self.slot_model = SequenceTagger.load('models/slot-model/final-model.pt')

    def get_intent(self, text: str):
        return predict_frame(self.frame_model, text.split(), 'frame')

    def get_slot(self, text: str):
        pred = predict_slot(self.slot_model, text.split(), 'slot')
        slots = []
        current_slot = None
        current_slot_value = []

        for frame in pred:
            slot = frame["slot"]
            if slot.startswith("B-"):
                if current_slot:
                    slots.append(Slot(name=current_slot, value=current_slot_value))
                current_slot = slot[2:]
                current_slot_value = [frame["form"]]
            elif slot.startswith("I-"):
                current_slot_value.append(frame["form"])

        if current_slot:
            slots.append(Slot(name=current_slot, value=current_slot_value))

        return slots

    def analyze(self, text: str):
        intent = self.get_intent(text)
        slots = self.get_slot(text)
        return Act(intent=intent, slots=slots)