✨ integrate constrained textgen

Signed-off-by: Peter <[email protected]>

app.py

@@ -85,6 +85,7 @@ def ask_gpt(
     top_p=0.95,
     top_k=25,
     temperature=0.5,
+    constrained_generation=True,
 ) -> str:
     """
     ask_gpt - helper function that asks the GPT model a question and returns the response
@@ -121,6 +122,7 @@ def ask_gpt(
         temperature=temperature,
         max_length=max_length,
         min_length=min_length,
+        constrained_generation  = constrained_generation,
     )
     gpt_et = time.perf_counter()
     gpt_rt = round(gpt_et - st, 2)

converse.py

@@ -10,6 +10,7 @@ import time
 
 from grammar_improve import remove_trailing_punctuation
 
+from constrained_generation import constrained_generation
 
 def discussion(
     prompt_text: str,
@@ -28,6 +29,7 @@ def discussion(
     num_return_sequences=1,
     device=-1,
     verbose=False,
+    constrained_generation=False,
 ):
     """
     discussion - a function that takes in a prompt and generates a response. This function is meant to be used in a conversation loop, and is the main function for the bot.
@@ -64,23 +66,46 @@ def discussion(
         pp.pprint(this_prompt, indent=4)
     # call the model
     print("\n... generating...")
-    bot_dialogue = gen_response(
-        this_prompt,
-        pipeline,
-        speaker,
-        responder,
-        timeout=timeout,
-        max_length=max_length,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        full_text=full_text,
-        no_repeat_ngram_size=no_repeat_ngram_size,
-        length_penalty=length_penalty,
-        num_return_sequences=num_return_sequences,
-        device=device,
-        verbose=verbose,
-    )
+    if constrained_generation:
+        response = constrained_generation(
+            prompt=this_prompt,
+            pipeline=pipeline,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            length_penalty=length_penalty,
+            repetition_penalty=1.0,
+            num_beams=4,
+            timeout=timeout,
+            verbose=verbose,
+            full_text=full_text,
+            speaker_name=speaker,
+            responder_name=responder,
+        )
+
+        bot_dialogue = consolidate_texts(
+            name_resp=responder,
+            model_resp=response,
+            name_spk=speaker,
+            verbose=verbose,
+            print_debug=True,
+        )
+    else:
+        bot_dialogue = gen_response(
+            this_prompt,
+            pipeline,
+            speaker,
+            responder,
+            timeout=timeout,
+            max_length=max_length,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
+            full_text=full_text,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            length_penalty=length_penalty,
+            num_return_sequences=num_return_sequences,
+            device=device,
+            verbose=verbose,
+        )
     if isinstance(bot_dialogue, list) and len(bot_dialogue) > 1:
         bot_resp = ", ".join(bot_dialogue)
     elif isinstance(bot_dialogue, list) and len(bot_dialogue) == 1: