Add a chat tab

gradio_app.py

@@ -1,5 +1,6 @@
-import random
 import os
+import random
+import string
 
 import gradio as gr
 import torch
@@ -9,6 +10,7 @@ import logging
 logger = logging.getLogger()
 logger.addHandler(logging.StreamHandler())
 
+DEBUG = os.environ.get("DEBUG", "false")[0] in "ty1"
 HF_AUTH_TOKEN = os.environ.get("HF_AUTH_TOKEN", None)
 DEVICE = os.environ.get("DEVICE", "cpu")  # cuda:0
 if DEVICE != "cpu" and not torch.cuda.is_available():
@@ -60,7 +62,10 @@ Este modelo ha sido entrenado con [Mesh Transformer JAX](https://github.com/king
 """
 
 FOOTER = """
-Para más información, visite el [repositorio del modelo](https://huggingface.co/bertin-project/bertin-gpt-j-6B).
+<div align=center>
+Para más información, visite el repositorio del modelo: <a href="https://huggingface.co/bertin-project/bertin-gpt-j-6B">BERTIN-GPT-J-6B</a>.
+<img src="https://visitor-badge.glitch.me/badge?page_id=spaces/bertin-project/bertin-gpt-j-6B"/>
+<div align=center>
 """.strip()
 
 EXAMPLES = [
@@ -74,6 +79,13 @@ Pregunta: ¿Quién cuidaba del hogar los dioses?
 Respuesta:""",
 ]
 
+AGENT = "BERTIN"
+USER = "ENTREVISTADOR"
+CONTEXT = """La siguiente conversación es un extracto de una entrevista a {AGENT} celebrada en Madrid para Radio Televisión Española:
+
+{USER}: Bienvenido, {AGENT}. Un placer tenerlo hoy con nosotros.
+{AGENT}: Gracias. El placer es mío."""
+
 class Normalizer:
     def remove_repetitions(self, text):
         """Remove repetitions"""
@@ -124,8 +136,6 @@ class TextGeneration:
     def generate(self, text, generation_kwargs):
         max_length = len(self.tokenizer(text)["input_ids"]) + generation_kwargs["max_length"]
         generation_kwargs["max_length"] = min(max_length, self.model.config.n_positions)
-        # generation_kwargs["num_return_sequences"] = 1
-        # generation_kwargs["return_full_text"] = False
         generated_text = None
         if text:
             for _ in range(10):
@@ -196,6 +206,64 @@ def expand_with_gpt(hidden, text, max_length, top_k, top_p, temperature, do_samp
     }
     return generator.generate(hidden or text, generation_kwargs)
 
+def chat_with_gpt(user, agent, context, user_message, history, max_length, top_k, top_p, temperature, do_sample, do_clean):
+    # agent = AGENT
+    # user = USER
+    generation_kwargs = {
+        "max_length": 25,
+        "top_k": top_k,
+        "top_p": top_p,
+        "temperature": temperature,
+        "do_sample": do_sample,
+        "do_clean": do_clean,
+        # "num_return_sequences": 1,
+        # "return_full_text": False,
+    }
+    message = user_message.split(" ", 1)[0].capitalize() + " " + user_message.split(" ", 1)[-1]
+    history = history or []  #[(f"{user}: Bienvenido. Encantado de tenerle con nosotros.", f"{agent}: Un placer, muchas gracias por la invitación.")]
+    context = context.format(USER=user or USER, AGENT=agent or AGENT).strip()
+    if context[-1] not in ".:":
+        context += "."
+    context_length = len(context.split())
+    history_take = 0
+    history_context = "\n".join(f"{user}: {history_message.capitalize()}.\n{agent}: {history_response}." for history_message, history_response in history[-len(history) + history_take:])
+    while len(history_context.split()) > generator.model.config.n_positions - (generation_kwargs["max_length"] + context_length):
+        history_take += 1
+        history_context = "\n".join(f"{user}: {history_message.capitalize()}.\n{agent}: {history_response}." for history_message, history_response in history[-len(history) + history_take:])
+        if history_take >= generator.model.config.n_positions:
+            break
+    context += history_context
+    for _ in range(5):
+        response = generator.generate(f"{context}\n\n{user}: {message}.\n", generation_kwargs)[1]
+        if DEBUG:
+            print("\n-----" + response + "-----\n")
+        response = response.split("\n")[-1]
+        if agent in response and response.split(agent)[-1]:
+            response = response.split(agent)[-1]
+        if user in response and response.split(user)[-1]:
+            response = response.split(user)[-1]
+        if response[0] in string.punctuation:
+            response = response[1:].strip()
+        if response.strip().startswith(f"{user}: {message}"):
+            response = response.strip().split(f"{user}: {message}")[-1]
+        if response.replace(".", "").strip() and message.replace(".", "").strip() != response.replace(".", "").strip():
+            break
+    if DEBUG:
+        print()
+        print("CONTEXT:")
+        print(context)
+        print()
+        print("MESSAGE")
+        print(message)
+        print()
+        print("RESPONSE:")
+        print(response)
+    if not response.strip():
+        response = random.choice(["No sé muy bien cómo contestar a eso.", "No estoy seguro.", "Prefiero no contestar.", "Ni idea.", "¿Podemos cambiar de tema?"])
+    history.append((user_message, response))
+    return history, history, ""
+
+
 with gr.Blocks() as demo:
     gr.Markdown(HEADER)
     with gr.Row():
@@ -247,24 +315,41 @@ with gr.Blocks() as demo:
                 # help="Si eliminar o no las palabras repetidas y recortar las últimas frases sin terminar.",
             )
         with gr.Column():
-            textbox = gr.Textbox(label="Texto", placeholder="Escriba algo (o seleccione un ejemplo) y pulse 'Generar'...", lines=8)
-            examples = gr.Dropdown(label="Ejemplos", choices=EXAMPLES, value=None, type="value")
-            hidden = gr.Textbox(visible=False, show_label=False)
-            with gr.Box():
-                # output = gr.Markdown()
-                output = gr.HighlightedText(label="Resultado", combine_adjacent=True, color_map={"BERTIN": "green", "ERROR": "red"})
-            with gr.Row():
-                btn = gr.Button("Generar")
-                btn.click(complete_with_gpt, inputs=[textbox, max_length, top_k, top_p, temperature, do_sample, do_clean], outputs=[textbox, hidden, output])
-                expand_btn = gr.Button("Añadir")
-                expand_btn.click(expand_with_gpt, inputs=[hidden, textbox, max_length, top_k, top_p, temperature, do_sample, do_clean], outputs=[textbox, hidden, output])
-
-                edit_btn = gr.Button("Editar", variant="secondary")
-                edit_btn.click(lambda x: (x, "", []), inputs=[hidden], outputs=[textbox, hidden, output])
-                clean_btn = gr.Button("Borrar", variant="secondary")
-                clean_btn.click(lambda: ("", "", [], ""), inputs=[], outputs=[textbox, hidden, output, examples])
-            examples.change(lambda x: x, inputs=[examples], outputs=[textbox])
+            with gr.Tabs():
+                with gr.TabItem("Generar"):
+                    textbox = gr.Textbox(label="Texto", placeholder="Escriba algo (o seleccione un ejemplo) y pulse 'Generar'...", lines=8)
+                    examples = gr.Dropdown(label="Ejemplos", choices=EXAMPLES, value=None, type="value")
+                    hidden = gr.Textbox(visible=False, show_label=False)
+                    with gr.Box():
+                        # output = gr.Markdown()
+                        output = gr.HighlightedText(label="Resultado", combine_adjacent=True, color_map={"BERTIN": "green", "ERROR": "red"})
+                    with gr.Row():
+                        generate_btn = gr.Button("Generar")
+                        generate_btn.click(complete_with_gpt, inputs=[textbox, max_length, top_k, top_p, temperature, do_sample, do_clean], outputs=[textbox, hidden, output])
+                        expand_btn = gr.Button("Añadir")
+                        expand_btn.click(expand_with_gpt, inputs=[hidden, textbox, max_length, top_k, top_p, temperature, do_sample, do_clean], outputs=[textbox, hidden, output])
+
+                        edit_btn = gr.Button("Editar", variant="secondary")
+                        edit_btn.click(lambda x: (x, "", []), inputs=[hidden], outputs=[textbox, hidden, output])
+                        clean_btn = gr.Button("Borrar", variant="secondary")
+                        clean_btn.click(lambda: ("", "", [], ""), inputs=[], outputs=[textbox, hidden, output, examples])
+                    examples.change(lambda x: x, inputs=[examples], outputs=[textbox])
+
+                with gr.TabItem("Charlar") as tab_chat:
+                    tab_chat.select(lambda: 25, inputs=[], outputs=[max_length])
+                    context = gr.Textbox(label="Contexto", value=CONTEXT, lines=5)
+                    with gr.Row():
+                        agent = gr.Textbox(label="Agente", value=AGENT)
+                        user = gr.Textbox(label="Usuario", value=USER)
+                    history = gr.Variable(default_value=[])
+                    chatbot = gr.Chatbot(color_map=("green", "gray"))
+                    with gr.Row():
+                        message = gr.Textbox(placeholder="Escriba aquí su mensaje y pulse 'Enviar'", show_label=False)
+                        chat_btn = gr.Button("Enviar")
+                    chat_btn.click(chat_with_gpt, inputs=[agent, user, context, message, history, max_length, top_k, top_p, temperature, do_sample, do_clean], outputs=[chatbot, history, message])
     gr.Markdown(FOOTER)
 
+
+
 demo.launch()
 # gr.Interface(complete_with_gpt, inputs=[textbox, max_length, top_k, top_p, temperature, do_sample, do_clean], outputs=[hidden, output]).launch()