From f4d9bff809feea6e33924523a8d226011b894fb7 Mon Sep 17 00:00:00 2001
From: kb <krzboj@st.amu.edu.pl>
Date: Mon, 10 Jun 2024 17:44:48 +0200
Subject: [PATCH] Dodaje opcje do mockowania NLU

---
 src/main.py                                  |  6 ++--
 src/service/natural_languag_understanding.py | 32 ++++++++++++--------
 2 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/src/main.py b/src/main.py
index 6e99efb..d8dd502 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,12 +1,12 @@
 from service.dialog_state_monitor import DialogStateMonitor
 from service.dialog_policy import DialogPolicy
 from service.natural_languag_understanding import NaturalLanguageUnderstanding
-from service.natural_language_generation import NaturalLanguageGeneration, parse_frame
+from service.natural_language_generation import NaturalLanguageGeneration
 from service.templates import templates
-from convlab.dialog_agent import PipelineAgent
 
 # initialize classes
-nlu = NaturalLanguageUnderstanding() # NLU
+
+nlu = NaturalLanguageUnderstanding(use_mocks=False) # NLU
 monitor = DialogStateMonitor() # DSM
 dialog_policy = DialogPolicy() # DP
 language_generation = NaturalLanguageGeneration(templates) # NLG
diff --git a/src/service/natural_languag_understanding.py b/src/service/natural_languag_understanding.py
index 0ca6c8d..dced5a0 100644
--- a/src/service/natural_languag_understanding.py
+++ b/src/service/natural_languag_understanding.py
@@ -1,4 +1,3 @@
-from flair.models import SequenceTagger
 from utils.nlu_utils import predict_single, predict_and_annotate
 from model.frame import Frame, Slot
 import random
@@ -42,7 +41,12 @@ SLOTS:
 """
 
 class NaturalLanguageUnderstanding():
-    def __init__(self):
+    def __init__(self, use_mocks=False):
+        if use_mocks:
+            self.use_mocks = True
+            return
+
+        from flair.models import SequenceTagger
         print("\n========================================================")
         print("Models are loading, it may take a moment, please wait...")
         print("========================================================\n")
@@ -86,14 +90,16 @@ class NaturalLanguageUnderstanding():
         return slots
 
     def predict(self, text: str):
-        act = self.__predict_intention(text)
-        slots = self.__predict_slot(text)
-        frame = Frame(source = 'user', act = act, slots = slots)
-        # uncomment to quickly mock the response
-        # frames = [
-        #     Frame(source="user", act = "inform/order", slots=[Slot(name="pizza", value="barcelona")]),
-        #     Frame(source="user", act = "welcomemsg", slots=[]),
-        #     Frame(source="user", act = "request/menu", slots=[]),
-        # ]
-        # return random.choice(frames)
-        return frame
\ No newline at end of file
+        if not self.use_mocks:
+            act = self.__predict_intention(text)
+            slots = self.__predict_slot(text)
+            frame = Frame(source = 'user', act = act, slots = slots)
+            return frame
+        else:
+            frames = [
+                Frame(source="user", act = "inform/order", slots=[Slot(name="pizza", value="barcelona")]),
+                Frame(source="user", act = "welcomemsg", slots=[]),
+                Frame(source="user", act = "request/menu", slots=[]),
+            ]
+            return random.choice(frames)
+        
\ No newline at end of file