Build a basic CLI app for the chatbot. Improve the way to read the prompt.

.gitignore

@@ -158,3 +158,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+.vscode/

model.py → chatbot.py

@@ -1,30 +1,41 @@
+"""
+Chatbot based on the notebook [How to generate text: using different decoding
+methods for language generation with
+Transformers](https://github.com/huggingface/blog/blob/main/notebooks/02_how_to_generate.ipynb)
+and the blog post [Create conversational agents using BLOOM:
+Part-1](https://medium.com/@fractal.ai/create-conversational-agents-using-bloom-part-1-63a66e6321c0).
+
+This code needs testing, as it is not fitted for a production model.
+
+It's a very basic chatbot that uses Causal Language Models from Transformers
+given an PROMPT.
+
+An example of a basic PROMPT is given in the file `prompt.txt` for a Spanish
+prompt.
+"""
+
+import argparse
+import torch
+
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from typing import Optional, Union
 
 
 class ChatBot:
     """
-    Chatbot based on the notebook [How to generate text: using different
-    decoding methods for language generation with
-    Transformers](https://github.com/huggingface/blog/blob/main/notebooks/02_how_to_generate.ipynb)
-    and the blog post [Create conversational agents using BLOOM:
-    Part-1](https://medium.com/@fractal.ai/create-conversational-agents-using-bloom-part-1-63a66e6321c0).
-
-    This code needs testing, as it is not fitted for a production model.
-
-    It's a very basic chatbot that uses Causal Language Models from Transformers given an PROMPT.
-
-    An example of a basic PROMPT is given by the BASE_PROMPT attribute.
+    Main class wrapper around the transformers models in order to build a basic
+    chatbot application.
 
     Parameters
     ----------
     base_model : str | AutoModelForCausalLM
         A name (path in hugging face hub) for a model, or the model itself.
     tokenizer : AutoTokenizer | None
-        Needed in case the base_model is a given model, otherwise it will load the same model
-        given by the base_model path.
+        Needed in case the base_model is a given model, otherwise it will load
+        the same model given by the base_model path.
     initial_prompt : str
-        A prompt for the model. Should follow the example given in `BASE_PROMPT`
+        A prompt for the model. Should follow the example given in
+        `BASE_PROMPT`
     keep_context : bool
         Whether to accumulate the context as the chatbot is used.
     creative : bool
@@ -34,23 +45,13 @@ class ChatBot:
     max_tokens : int
         Max number of tokens to generate in the chat.
     human_identifier : str
-        The string that will identify the human speaker in the prompt (e.g. HUMAN).
+        The string that will identify the human speaker in the prompt (e.g.
+        HUMAN).
     bot_identifier : str
-        The string that will identify the bot speaker in the prompt (e.g. EXPERT).
+        The string that will identify the bot speaker in the prompt (e.g.
+        EXPERT).
     """
 
-    BASE_PROMPT = """
-    The following is a conversation with a movie EXPERT.
-    The EXPERT helps the HUMAN define their personal preferences and provide
-    multiple options to select from, it also helps in selecting the best option.
-    The EXPERT is conversational, optimistic, flexible, empathetic, creative and
-    humanly in generating responses.
-
-    HUMAN: Hello, how are you?
-    EXPERT: Fine, thanks. I am here to help you by recommending movies.
-    """.strip()
-
-
     def __init__(self,
                  base_model: Union[str, AutoModelForCausalLM],
                  tokenizer: Optional[AutoTokenizer] = None,
@@ -58,8 +59,8 @@ class ChatBot:
                  keep_context: bool = False,
                  creative: bool = False,
                  max_tokens: int = 50,
-                 human_identifier: str = "HUMAN",
-                 bot_identifier: str = "EXPERT"):
+                 human_identifier: str = 'HUMAN',
+                 bot_identifier: str = 'EXPERT'):
         if isinstance(base_model, str):
             self.model = AutoModelForCausalLM.from_pretrained(
                 base_model,
@@ -73,7 +74,12 @@ class ChatBot:
             self.model = base_model
             self.tokenizer = tokenizer
 
-        self.initial_prompt = initial_prompt if initial_prompt is not None else self.BASE_PROMPT
+        if initial_prompt is None:
+            with open('./prompt.txt', 'r') as fh:
+                self.initial_prompt = fh.read()
+        else:
+            self.initial_prompt = initial_prompt
+
         self.keep_context = keep_context
         self.context = ''
         self.creative = creative
@@ -85,11 +91,12 @@ class ChatBot:
         """
         Generates a response from the prompt (and optionally the context) where
         it adds the `input_text` as if it was part of the HUMAN dialog
-        (identified by `self.human_identifier`), and prompts the bot (identified
-        by `self.bot_identifier`) for a response. As the bot might continue the
-        conversation beyond the scope, it trims the output so it only shows the
-        first dialog given by the bot, following the idea presented in the
-        Medium blog post for creating conversational agents (link above).
+        (identified by `self.human_identifier`), and prompts the bot
+        (identified by `self.bot_identifier`) for a response. As the bot might
+        continue the conversation beyond the scope, it trims the output so it
+        only shows the first dialog given by the bot, following the idea
+        presented in the Medium blog post for creating conversational agents
+        (link above).
 
         Parameters
         ----------
@@ -101,31 +108,94 @@ class ChatBot:
         str
             The output given by the bot, trimmed for better control.
         """
+        # Setup the prompt given the initial prompt and add the words that
+        # start the dialog between the human and the bot. Give space for the
+        # model to continue from the prompt
         prompt = self.initial_prompt + self.context
         prompt += f'{self.human_identifier}: {input_text}\n'
-        prompt += f'{self.bot_identifier}: '
+        prompt += f'{self.bot_identifier}: '  # check the space after the colon
 
         input_ids = self.tokenizer.encode(prompt, return_tensors='pt')
         if self.creative:
+            # In case you want the bot to be creative, we sample using `top_k`
+            # and `top_p`
             output = self.model.generate(
                 input_ids,
                 do_sample=True,
                 max_length=input_ids.shape[1] + self.max_tokens,
                 top_k=50,
-                top_p=0.95,
-                num_return_sequences=1
+                top_p=0.95
             )[0]
         else:
+            # Otherwise we return the most probable token
             output = self.model.generate(
                 input_ids,
                 max_length=input_ids.shape[1] + self.max_tokens
             )[0]
 
+        # Decode the output, removing special tokens for the model (like
+        # `[CLS]` and similar)
         decoded_output = self.tokenizer.decode(output, skip_special_tokens=True)
+
+        # Trim the output, first by removing the original prompt
         trimmed_output = decoded_output[len(prompt):]
+
+        # Then we find the stop token, in this case the human identifier, and
+        # we get up to that point
         trimmed_output = trimmed_output[:trimmed_output.find(f'{self.human_identifier}:')]
 
         if self.keep_context:
-            self.context += trimmed_output
-
-        return trimmed_output.strip()
+            # If we want to keep the context of the conversation we add the
+            # trimmed output so far
+            self.context += prompt + trimmed_output
+
+        return trimmed_output.strip()  # we only return the trimmed output
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model-name', '-m',
+                        default='bigscience/bloom-560m',
+                        help="Name of the base model to use for the chatbot")
+    parser.add_argument('--prompt', '-p',
+                        default='./prompt.txt',
+                        help="Path to the file with the prompt to use")
+    parser.add_argument('--keep-context', '-k',
+                        action='store_true',
+                        help="Keep context of the conversation.")
+    parser.add_argument('--creative', '-c',
+                        action='store_true',
+                        help="Make the bot creative when answering.")
+    parser.add_argument('--random-seed', '-r',
+                        default=42,
+                        help="Seed number for the creative bot.",
+                        type=int)
+    parser.add_argument('--human-identifier', '-i',
+                        default='HUMANO',
+                        help="Name of the human identifier.")
+    parser.add_argument('--bot-identifier', '-b',
+                        default='EXPERTO',
+                        help="Name of the bot identifier.")
+
+    args = parser.parse_args()
+
+    torch.manual_seed(args.random_seed)
+
+    with open(args.prompt, 'r') as fh:
+        initial_prompt = fh.read()
+
+    chatbot = ChatBot(
+        base_model=args.model_name,
+        initial_prompt=initial_prompt,
+        keep_context=args.keep_context,
+        creative=args.creative,
+        human_identifier=args.human_identifier,
+        bot_identifier=args.bot_identifier
+    )
+
+    print("Write `exit` or `quit` to quit")
+    while True:
+        input_text = input('> ')
+        if input_text == 'exit' or input_text == 'quit':
+            break
+        print(chatbot.chat(input_text))

prompt.txt

@@ -0,0 +1,6 @@
+La siguiente es una conversación entre un HUMANO y un bot EXPERTO en software libre.
+El EXPERTO le ayuda al HUMANO con preguntas acerca de software libre.
+El EXPERTO es conversacional, optimista, flexible, creativo y genera respuestas parecidas a un humano.
+
+HUMANO: Hola, ¿Cómo estás?
+EXPERTO: Hola, muy bien. Estoy acá para ayudarte con preguntas respecto al software libre.