From 7a5064f19e00e81c73552b733c69877f09b435c0 Mon Sep 17 00:00:00 2001
From: Adrian Klessa <50918271+AdrianKlessa@users.noreply.github.com>
Date: Mon, 10 Jun 2024 20:48:24 +0200
Subject: [PATCH] Restart + "clean" console command after user is finished

---
 dialog_with_nlg.py | 38 +++++++++++++++++++++++---------------
 1 file changed, 23 insertions(+), 15 deletions(-)

diff --git a/dialog_with_nlg.py b/dialog_with_nlg.py
index fc13fe1..df7164b 100644
--- a/dialog_with_nlg.py
+++ b/dialog_with_nlg.py
@@ -49,11 +49,11 @@ class Model:
 
         r = self.nlu(msg)
         slots = r['slots']
-        #print(r)
+        # print(r)
         r = self.dst(r)
-        #print(r)
+        # print(r)
         r = self.dp()
-        #print(r)
+        # print(r)
         r = self.nlg(r, slots)
         if debug:
             print(r)
@@ -193,7 +193,7 @@ class NLG():
     def __init__(self, state):
         self.model = AutoModelForSeq2SeqLM.from_pretrained("filnow/nlg-umt5-pol")
         self.tokenizer = AutoTokenizer.from_pretrained("google/umt5-small")
-        self.nlg_pipeline  = pipeline('summarization', model=self.model, tokenizer=self.tokenizer)
+        self.nlg_pipeline = pipeline('summarization', model=self.model, tokenizer=self.tokenizer)
 
         self.messages = {
             "welcomemsg": [
@@ -231,31 +231,32 @@ class NLG():
     def __call__(self, act: str, slots: List[Tuple[str, str]]) -> str:
         if act == 'welcomemsg':
             return random.choice(self.messages["welcomemsg"])
-        
+
         elif act in ["inform", "request", "select"]:
             if slots == []:
                 return "Przepraszam nie rozumiem. Podaj więcej informacji."
             else:
                 text = [f"{slot[0]}[{slot[1]}]" for slot in slots if slot[1] is not None]
                 return self.nlg_pipeline(f'generate text: {", ".join(text)}')[0]['summary_text']
-        
+
         elif act == "canthelp.missing_slot_value" or act == "canthelp":
             return random.choice(self.messages["canthelp"])
-            
+
         elif act == "bye":
             return random.choice(self.messages["bye"])
-        
+
         elif act == 'affirm':
             return random.choice(self.messages["affirm"])
-        
+
         elif act == "repeat":
             return random.choice(self.messages["repeat"])
-        
+
         elif act == "reqmore":
             return random.choice(self.messages["reqmore"])
-        
+
         elif act == "offer":
-            return "Proszę oto menu zeskanuj kod QR aby je zobaczyć." 
+            return "Proszę oto menu zeskanuj kod QR aby je zobaczyć."
+
 
 def dialogue_test():
     model = Model()
@@ -285,12 +286,12 @@ def dialogue_test():
     response = model("uniwersytetu poznanskiego 4 61-614 poznan")
     print()
     # jezeli sprobuje dokonac zamowienia bez podania potrzebnych informacji prosimy o nie
-    #response = model("Dobrze, nie mogę się już doczekać.")
+    # response = model("Dobrze, nie mogę się już doczekać.")
     response = model("Super, to zatem wszystko!")
     print()
     # jezeli wybierze rodzaj platnosci to zapisz wybor i poinformuj o nim
     # response = model("karta")
-    #response = model("Poproszę blikiem z góry")
+    # response = model("Poproszę blikiem z góry")
     response = model("Zapłacę kartą przy odbiorze")
     print()
     # jezeli potwiedzi zamowienie to zakoncz zamawianie sukcesem i wypisz calosc
@@ -325,10 +326,17 @@ def dialogue_test():
     response = model("Dobrze, nie mogę się już doczekać.")
     print()
 
+
 if __name__ == "__main__":
     model = Model()
     print("Chatbot Jarvis\n--------------")
     while True:
         print("\nUżytkownik: ")
         user_input = input()
-        response = model(user_input, debug=False)
\ No newline at end of file
+        while user_input == "SYSTEM_FINISH":
+            print("\n\n\n\n\n\n\n\n\n\n\n\n")
+            model = Model()
+            print("Chatbot Jarvis\n--------------")
+            print("\nUżytkownik: ")
+            user_input = input()
+        response = model(user_input, debug=False)