Update demo

gradio_app.py

@@ -50,9 +50,10 @@ LOW_CPU_MEM = bool(os.environ.get("LOW_CPU_MEM", False if DEVICE == "cpu" else T
 MODEL_NAME = os.environ.get("MODEL_NAME", "bertin-project/bertin-gpt-j-6B")
 MODEL_REVISION = os.environ.get("MODEL_REVISION", "main")
 MAX_LENGTH = int(os.environ.get("MAX_LENGTH", 1024))
-HEADER_INFO = """
-# BERTIN GPT-J-6B
-Spanish BERTIN GPT-J-6B Model.
+display_model_name = "BERTIN GPT-J-6B" if MODEL_NAME == "bertin-project/bertin-gpt-j-6B" else MODEL_NAME.upper()
+HEADER_INFO = f"""
+# {display_model_name}
+Spanish {display_model_name} Model.
 """.strip()
 LOGO = "https://huggingface.co/bertin-project/bertin-roberta-base-spanish/resolve/main/images/bertin.png"
 HEADER = f"""
@@ -83,7 +84,7 @@ span.generated-text {{
 <div align=center>
 <img src="{LOGO}" width=150/>
 
-# BERTIN GPT-J-6B
+# {display_model_name}
 
 BERTIN proporciona una serie de modelos de lenguaje en Español entrenados en abierto.
 
@@ -92,10 +93,10 @@ Este modelo ha sido entrenado con [Mesh Transformer JAX](https://github.com/king
 </div>
 """
 
-FOOTER = """
+FOOTER = f"""
 <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"/>
+Para más información, visite el repositorio del modelo: <a href="https://huggingface.co/{MODEL_NAME}">{display_model_name}</a>.
+<img src="https://visitor-badge.glitch.me/badge?page_id=spaces/{MODEL_NAME}"/>
 <div align=center>
 """.strip()
 
@@ -103,6 +104,7 @@ EXAMPLES = [
     "",
     "Érase una vez,",
     "¿Cuál es la capital de Francia? Respuesta:",
+    "En un lugar de la Mancha, de cuyo nombre no quiero acordarme, no ha mucho tiempo que vivía un hidalgo de los de lanza en astillero, adarga antigua, rocín flaco y galgo corredor.",
     """Los templos egipcios fueron construidos para el culto oficial de los dioses y la conmemoración de los faraones del Antiguo Egipto en las regiones bajo su dominio. Los templos eran vistos como el hogar de los dioses o faraones deificados a quienes eran dedicados, y en ellos los faraones y el clero egipcio llevaban a cabo diversos rituales, las funciones centrales de la religión egipcia: realizar ofrendas a sus dioses, recrear pasajes mitológicos mediante festivales y protegerse de las fuerzas del caos. Estos rituales eran vistos como necesarios para que los dioses mantuvieran la maat, el orden divino del universo.
 
 El cuidado del hogar de los dioses era obligación de los faraones, que dedicaron ingentes cantidades de recursos para la construcción y el mantenimiento de los templos. Por necesidad, los faraones delegaban la mayoría de los rituales en una amplia casta sacerdotal, aunque la mayor parte del pueblo llano permanecía al margen de la participación directa en las ceremonias por tener prohibido el acceso a las zonas más sagradas de los templos. A pesar de ello, el templo siempre fue un importante centro religioso para todos los egipcios, que iban a ellos a rezar, realizar ofrendas y buscar la guía de los oráculos.
@@ -111,7 +113,8 @@ Pregunta: ¿Quién cuidaba del hogar los dioses?
 Respuesta:""",
 ]
 
-AGENT = "BERTIN"
+AGENT = os.environ.get("AGENT_NAME", "BERTIN") 
+PREV = "PREV"
 USER = "ENTREVISTADOR"
 CONTEXT = """La siguiente conversación es un extracto de una entrevista a {AGENT} celebrada en Madrid para Radio Televisión Española:
 
@@ -148,6 +151,9 @@ class TextGeneration:
         self.tokenizer = AutoTokenizer.from_pretrained(
             self.model_name_or_path, revision=MODEL_REVISION, use_auth_token=HF_AUTH_TOKEN if HF_AUTH_TOKEN else None,
         )
+        self.tokenizer_prefix_space = AutoTokenizer.from_pretrained(
+            self.model_name_or_path, add_prefix_space=True, revision=MODEL_REVISION, use_auth_token=HF_AUTH_TOKEN if HF_AUTH_TOKEN else None,
+        )
         self.model = AutoModelForCausalLM.from_pretrained(
             self.model_name_or_path, revision=MODEL_REVISION,
             use_auth_token=HF_AUTH_TOKEN if HF_AUTH_TOKEN else None,
@@ -167,24 +173,40 @@ class TextGeneration:
 
 
     def generate(self, text, generation_kwargs, previous_text=None):
+        do_clean = generation_kwargs.pop("do_clean", False)
+        bad_words = generation_kwargs.pop("bad_words", "")
+        if bad_words:
+            generation_kwargs["bad_words_ids"] = self.tokenizer_prefix_space(
+                [word.strip() for word in bad_words.split(",")], add_special_tokens=False
+            ).input_ids
+        if "repetition_penalty" in generation_kwargs:
+            generation_kwargs["repetition_penalty"] = float(generation_kwargs["repetition_penalty"])
         input_text = previous_text or text
-        max_length = len(self.tokenizer(input_text)["input_ids"]) + generation_kwargs["max_length"]
-        generation_kwargs["max_length"] = min(max_length, self.model.config.n_positions)
+        # max_length = len(self.tokenizer(input_text)["input_ids"]) + generation_kwargs["max_length"]
+        # generation_kwargs["max_length"] = min(max_length, self.model.config.n_positions)
+        generation_kwargs["max_new_tokens"] = generation_kwargs.pop("max_length", 50)
         generated_text = None
         if input_text:
+            pre_input_text = ""
+            input_ids = self.tokenizer(input_text).input_ids
+            if len(input_ids) + generation_kwargs["max_new_tokens"] >= 2048:
+                prompt_cutoff = 2048 - generation_kwargs["max_new_tokens"] + 1
+                pre_input_text = self.tokenizer.decode(input_ids[:-prompt_cutoff])
+                input_text = self.tokenizer.decode(input_ids[-prompt_cutoff:])
             for _ in range(10):
-                generated_text = self.generator(
+                generated_text = pre_input_text + (" " if do_clean else "") + self.generator(
                     input_text,
                     **generation_kwargs,
-                )[0]["generated_text"]
+                )[0]["generated_text"]                
+                input_text = self.tokenizer.decode(input_ids)
                 if generated_text.strip().startswith(input_text):
                     generated_text = generated_text.replace(input_text, "", 1).strip()
-                if generation_kwargs["do_clean"]:
+                if do_clean:
                     generated_text = cleaner.clean_txt(generated_text)
                 if generated_text:
                     if previous_text and previous_text != text:
                         diff = [
-                            (text, None), (previous_text.replace(text, " ", 1).strip(), " "), (generated_text, AGENT)
+                            (text, None), (previous_text.replace(text, " ", 1).strip(), PREV), (generated_text, AGENT)
                         ]
                     else:
                         diff = [(text, None), (generated_text, AGENT)]
@@ -195,7 +217,7 @@ class TextGeneration:
             if not generated_text:
                 return (
                     "",
-                    [("Tras 10 intentos BERTIN no generó nada. Pruebe cambiando las opciones.", "ERROR")]
+                    [(f"Tras 10 intentos {AGENT} no generó nada. Pruebe cambiando las opciones.", "ERROR")]
                 )
         return (
             "",
@@ -215,36 +237,51 @@ cleaner = Normalizer()
 generator = load_text_generator()
 
 
-def complete_with_gpt(text, max_length, top_k, top_p, temperature, do_sample, do_clean):
+def complete_with_gpt(text, max_length, top_k, top_p, penalty_alpha, num_beams, temperature, repetition_penalty, no_repeat_ngram_size, bad_words, do_sample, do_clean):
     generation_kwargs = {
         "max_length": max_length,
         "top_k": top_k,
         "top_p": top_p,
+        "penalty_alpha": penalty_alpha,
+        "num_beams": num_beams,
         "temperature": temperature,
+        "repetition_penalty": repetition_penalty,
+        "no_repeat_ngram_size": no_repeat_ngram_size,
+        "bad_words": bad_words,
         "do_sample": do_sample,
         "do_clean": do_clean,
     }
     return generator.generate(text, generation_kwargs)
 
-def expand_with_gpt(hidden, text, max_length, top_k, top_p, temperature, do_sample, do_clean):
+def expand_with_gpt(hidden, text, max_length, top_k, top_p, penalty_alpha, num_beams, temperature, repetition_penalty, no_repeat_ngram_size, bad_words, do_sample, do_clean):
     generation_kwargs = {
         "max_length": max_length,
         "top_k": top_k,
         "top_p": top_p,
+        "penalty_alpha": penalty_alpha,
+        "num_beams": num_beams,
         "temperature": temperature,
+        "repetition_penalty": repetition_penalty,
+        "no_repeat_ngram_size": no_repeat_ngram_size,
+        "bad_words": bad_words,
         "do_sample": do_sample,
         "do_clean": do_clean,
     }
     return generator.generate(text, generation_kwargs, previous_text=hidden)
 
-def chat_with_gpt(agent, user, context, user_message, history, max_length, top_k, top_p, temperature, do_sample, do_clean):
+def chat_with_gpt(agent, user, context, user_message, history, max_length, top_k, top_p, penalty_alpha, num_beams, temperature, repetition_penalty, no_repeat_ngram_size, bad_words, do_sample, do_clean):
     # agent = AGENT
     # user = USER
     generation_kwargs = {
-        "max_length": 25,
+        "max_length": max_length,
         "top_k": top_k,
         "top_p": top_p,
+        "penalty_alpha": penalty_alpha,
+        "num_beams": num_beams,
         "temperature": temperature,
+        "repetition_penalty": repetition_penalty,
+        "no_repeat_ngram_size": no_repeat_ngram_size,
+        "bad_words": bad_words,
         "do_sample": do_sample,
         "do_clean": do_clean,
         # "num_return_sequences": 1,
@@ -295,7 +332,7 @@ def chat_with_gpt(agent, user, context, user_message, history, max_length, top_k
         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?"])
+        response = random.choice(["No sé muy bien cómo contestar a eso.", "No puedo contestar con seguridad.", "Prefiero no contestar.", "Ni idea.", "¿Podemos cambiar de tema?"])
     history.append((user_message, response))
     return history, history, ""
 
@@ -304,54 +341,103 @@ def chat_with_gpt(agent, user, context, user_message, history, max_length, top_k
 with gr.Blocks() as demo:
     gr.Markdown(HEADER)
     with gr.Row():
-        with gr.Group():
-            with gr.Box():
-                gr.Markdown("Opciones")
-            max_length = gr.Slider(
-                label='Longitud máxima',
-                # help="Número máximo (aproximado) de palabras a generar.",
-                minimum=1,
-                maximum=MAX_LENGTH,
-                value=50,
-                step=1
-            )
-            top_k = gr.Slider(
-                label='Top-k',
-                # help="Número de palabras con alta probabilidad a mantener para el filtrado `top-k`",
-                minimum=40,
-                maximum=80,
-                value=50,
-                step=1
-            )
-            top_p = gr.Slider(
-                label='Top-p',
-                # help="Solo las palabras más probables con probabilidades que sumen `top_p` o más se mantienen para la generación.",
-                minimum=0.0,
-                maximum=1.0,
-                value=0.95,
-                step=0.01
-            )
-            temperature = gr.Slider(
-                label='Temperatura',
-                # help="Valor utilizado para modular las probabilidades de las siguientes palabras generadas.",
-                minimum=0.1,
-                maximum=10.0,
-                value=0.8,
-                step=0.05
-            )
-            do_sample = gr.Checkbox(
-                label='¿Muestrear?',
-                value = True,
-                # options=(True, False),
-                # help="Si no se muestrea se usará una decodificación voraz (_greedy_).",
-            )
-            do_clean = gr.Checkbox(
-                label='¿Limpiar texto?',
-                value = True,
-                # options=(True, False),
-                # help="Si eliminar o no las palabras repetidas y recortar las últimas frases sin terminar.",
-            )
-        with gr.Column():
+        with gr.Column(scale=1):
+            with gr.Group():
+                with gr.Box():
+                    gr.Markdown("Opciones")
+                    with gr.Tabs():
+                        with gr.TabItem("Generación"):
+                            max_length = gr.Slider(
+                                label='Palabras a generar',
+                                # help="Número máximo (aproximado) de palabras a generar.",
+                                minimum=1,
+                                maximum=MAX_LENGTH,
+                                value=50,
+                                step=1
+                            )
+                            top_k = gr.Slider(
+                                label='Top-k',
+                                # help="Número de palabras con alta probabilidad a mantener para el filtrado `top-k`",
+                                minimum=0,
+                                maximum=80,
+                                value=50,
+                                step=1
+                            )
+                            top_p = gr.Slider(
+                                label='Top-p',
+                                # help="Solo las palabras más probables con probabilidades que sumen `top_p` o más se mantienen para la generación.",
+                                minimum=0.01,
+                                maximum=5.0,
+                                value=0.95,
+                                step=0.01
+                            )
+                            penalty_alpha = gr.Slider(
+                                label='Penalización (alpha)',
+                                # help="Penalización para contrastive search.",
+                                minimum=0.0,
+                                maximum=1.0,
+                                value=0.0,
+                                step=0.01
+                            )
+                            num_beams = gr.Slider(
+                                label='Haces (beams)',
+                                # help="Número de beams para búsqueda.",
+                                minimum=1,
+                                maximum=50,
+                                value=1,
+                                step=1
+                            )
+                            temperature = gr.Slider(
+                                label='Temperatura',
+                                # help="Valor utilizado para modular las probabilidades de las siguientes palabras generadas.",
+                                minimum=0.0,
+                                maximum=10.0,
+                                value=0.8,
+                                step=0.05
+                            )
+                            do_sample = gr.Checkbox(
+                                label='¿Muestrear?',
+                                value = True,
+                                # options=(True, False),
+                                # help="Si no se muestrea se usará una decodificación voraz (_greedy_).",
+                            )
+                            do_clean = gr.Checkbox(
+                                label='¿Limpiar texto?',
+                                value = False,
+                                # options=(True, False),
+                                # help="Si eliminar o no las palabras repetidas y recortar las últimas frases sin terminar.",
+                            )
+                        with gr.TabItem("Control de repetición"):
+                            repetition_penalty = gr.Slider(
+                                label='Penalización por repetición',
+                                help="Un valor de 1 significa no penalización.",
+                                minimum=1.0,
+                                maximum=10.0,
+                                value=1.0,
+                                step=0.01
+                            )
+                            no_repeat_ngram_size = gr.Slider(
+                                label='No repetir ngrams de tamaño',
+                                minimum=0,
+                                maximum=10,
+                                value=0,
+                                step=1
+                            )
+                            bad_words = gr.Textbox(
+                                label="Palabras a evitar",
+                                info="Lista de palabras separadas por comas",
+                                lines=1,
+                                value="",
+                            )
+                    with gr.Accordion("Estrategias", open=False):
+                        gr.Markdown("""
+                        - **greedy decoding** si `num_beams=1` y `do_sample=False`
+                        - **contrastive search** si `penalty_alpha>0.0` y `top_k>1`
+                        - **multinomial sampling** si `num_beams=1` y `do_sample=True`
+                        - **beam-search decoding** si `num_beams>1` y `do_sample=False`
+                        - **beam-search multinomial sampling** si `num_beams>1` y `do_sample=True`
+                        """)
+        with gr.Column(scale=4):
             with gr.Tabs():
                 with gr.TabItem("Generar"):
                     textbox = gr.Textbox(label="Texto", placeholder="Escriba algo (o seleccione un ejemplo) y pulse 'Generar'...", lines=8)
@@ -363,13 +449,14 @@ with gr.Blocks() as demo:
                             elem_id="htext",
                             label="Resultado",
                             combine_adjacent=True,
-                            color_map={AGENT: "green", "ERROR": "red", " ": "blue"},
+                        ).style(
+                            color_map={AGENT: "green", "ERROR": "red", PREV: "blue"},
                         )
                     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=[hidden, output])
+                        generate_btn.click(complete_with_gpt, inputs=[textbox, max_length, top_k, top_p, penalty_alpha, num_beams, temperature, repetition_penalty, no_repeat_ngram_size, bad_words, do_sample, do_clean], outputs=[hidden, output], api_name="generate")
                         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=[hidden, output])
+                        expand_btn.click(expand_with_gpt, inputs=[hidden, textbox, max_length, top_k, top_p, penalty_alpha, num_beams, temperature, repetition_penalty, no_repeat_ngram_size, bad_words, do_sample, do_clean], outputs=[hidden, output])
 
                         edit_btn = gr.Button("Editar", variant="secondary")
                         edit_btn.click(lambda x: (x, "", []), inputs=[hidden], outputs=[textbox, hidden, output])
@@ -378,23 +465,23 @@ with gr.Blocks() as demo:
                     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])
+                    # 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(value=[])
-                    chatbot = gr.Chatbot(color_map=("green", "gray"))
+                    chatbot = gr.Chatbot().style(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])
+                    chat_btn.click(chat_with_gpt, inputs=[agent, user, context, message, history, max_length, top_k, top_p, penalty_alpha, num_beams, temperature, repetition_penalty, no_repeat_ngram_size, bad_words, do_sample, do_clean], outputs=[chatbot, history, message])
     gr.Markdown(FOOTER)
 
-with gr.Interface(lambda: None, inputs=["text", max_length, top_k, top_p, temperature, do_sample, do_clean], outputs=[hidden, output]) as iface:
-    demo.examples = None
-    demo.predict_durations = []
-    demo.input_components = iface.input_components
-    demo.output_components = iface.output_components
-
-demo.launch()
+# with gr.Interface(lambda: None, inputs=["text", max_length, top_k, top_p, penalty_alpha, num_beams, temperature, do_sample, do_clean], outputs=[hidden, output]) as iface:
+#     demo.examples = None
+#     demo.predict_durations = []
+#     demo.input_components = iface.input_components
+#     demo.output_components = iface.output_components
+demo.queue()
+demo.launch(share=True)