Duplicate from joaogante/contrastive_search_generation

Co-authored-by: Joao Gante <[email protected]>

.gitattributes

@@ -0,0 +1,31 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: Constrastive Search Generation
+emoji: 🐠
+colorFrom: indigo
+colorTo: pink
+sdk: gradio
+sdk_version: 3.5
+app_file: app.py
+pinned: false
+license: mit
+duplicated_from: joaogante/contrastive_search_generation
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,204 @@
+import time
+from functools import lru_cache
+
+import torch
+import gradio as gr
+from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
+
+
+@lru_cache(maxsize=1)  # only cache the latest model
+def get_model_and_tokenizer(model_id):
+    config = AutoConfig.from_pretrained(model_id)
+    if config.is_encoder_decoder:
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(model_id)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    return model, tokenizer
+
+
+@lru_cache(maxsize=32768)  # cache up to 32k examples
+def run_generation(
+    text,
+    model_id,
+    max_new_tokens,
+    alpha=0.0,
+    top_k=0,
+    num_beams=1,
+    do_sample=False,
+    top_p=0.0,
+    seed=0
+):
+    model, tokenizer = get_model_and_tokenizer(model_id)
+
+    inputs = tokenizer(text, return_tensors='pt')
+    if seed:
+        torch.manual_seed(seed)
+
+    start = time.time_ns()
+    contrastive_ids = model.generate(
+        # from the tokenizer
+        **inputs,
+        # fixed arguments
+        num_return_sequences=1,
+        early_stopping=True,
+        # variable arguments
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        num_beams=num_beams,
+        penalty_alpha=alpha or None,
+        top_k=top_k or None,
+        top_p=top_p or None,
+    )
+    end = time.time_ns()
+
+    contrastive_time = (end - start) / 1e6
+    contrastive_text = tokenizer.decode(contrastive_ids[0], skip_special_tokens=True)
+    return contrastive_text, contrastive_time
+
+
+def generate_beam_search(text, model_id, max_new_tokens, alpha, k, num_beams):
+    contrastive_text, contrastive_time = run_generation(text, model_id, max_new_tokens, alpha=alpha, top_k=k)
+    beam_search_text, beam_search_time = run_generation(text, model_id, max_new_tokens, num_beams=num_beams)
+    return contrastive_text, contrastive_time, beam_search_text, beam_search_time
+
+
+def generate_top_k(text, model_id, max_new_tokens, alpha, k, top_k, seed):
+    contrastive_text, contrastive_time = run_generation(text, model_id, max_new_tokens, alpha=alpha, top_k=k)
+    top_k_text, top_k_time = run_generation(
+        text, model_id, max_new_tokens, top_k=top_k, seed=seed, do_sample=True
+    )
+    return contrastive_text, contrastive_time, top_k_text, top_k_time
+
+
+def generate_top_p(text, model_id, max_new_tokens, alpha, k, top_p, seed):
+    contrastive_text, contrastive_time = run_generation(text, model_id, max_new_tokens, alpha=alpha, top_k=k)
+    top_p_text, top_p_time = run_generation(
+        text, model_id, max_new_tokens, top_p=top_p, seed=seed, do_sample=True
+    )
+    return contrastive_text, contrastive_time, top_p_text, top_p_time
+
+
+demo = gr.Blocks()
+
+with demo:
+    gr.Markdown(
+        """
+        # Contrastive Search Generation comparison
+
+        Credits to the contrastive search generation [paper](https://arxiv.org/abs/2202.06417) authors, including
+        @[pangpang666](https://huggingface.co/pangpang666) and @[GMFTBY](https://huggingface.co/GMFTBY). Check out the
+        follow-up [work](https://arxiv.org/abs/2210.14140), which demonstrates the usefulness of the technique with
+        off-the-shelf LLMs, as well as their [HF guest blog post](https://huggingface.co/blog/introducing-csearch).
+
+        From the paper:
+        "At each decoding step, the key ideas of contrastive search are (i) the generated output should be selected
+        from the set of most probable candidates predicted by the model; and (ii) the generated output should be
+        discriminative enough with respect to the previous context. In this way, the generated text can (i) better
+        maintain the semantic coherence with respect to the prefix while (ii) avoiding model degeneration."
+
+        🚨 Warnings: 🚨
+        - Avoid using large models (> 1GB) in this demo. It will take a long time to load the model and generate text.
+        - Too slow/long queue? Check our
+        [colab](https://colab.research.google.com/github/huggingface/blog/blob/main/notebooks/115_introducing_contrastive_search.ipynb)
+        instead.
+        """
+    )
+    with gr.Tabs():
+        with gr.TabItem("vs. Beam Search"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("## Inputs ✍️")
+                    gr.Markdown("General options:")
+                    model_id = gr.Text(value="facebook/opt-125m", label="Model Repository")
+                    input_text = gr.Textbox(value="DeepMind Company is", lines=5, label="Input Text")
+                    max_new_tokens = gr.Slider(value=50, minimum=1, maximum=256, label="New tokens to generate")
+                    gr.Markdown("Contrastive Search options:")
+                    alpha = gr.Slider(value=0.6, minimum=0.01, maximum=1.0, step=0.01, label="Alpha")
+                    k = gr.Slider(value=6, minimum=1, maximum=20, step=1, label="K")
+                    gr.Markdown("Beam Search options:")
+                    num_beams = gr.Slider(value=4, minimum=1, maximum=16, step=1, label="Number of beams")
+                    generate_button = gr.Button(value="Generate", label="Generate")
+
+                with gr.Column():
+                    gr.Markdown("## Outputs 🤖")
+                    gr.Markdown("Contrastive Search generation:")
+                    text_contrastive = gr.Textbox(value="", label="")
+                    time_contrastive = gr.Number(value=0.0, precision=1, label="Generation time (ms)")
+                    gr.Markdown("Beam Search generation:")
+                    text_beam_search = gr.Textbox(value="", label="")
+                    time_beam_search = gr.Number(value=0.0, precision=1, label="Generation time (ms)")
+
+            # actions
+            generate_button.click(
+                fn=generate_beam_search,
+                inputs=[input_text, model_id, max_new_tokens, alpha, k, num_beams],
+                outputs=[text_contrastive, time_contrastive, text_beam_search, time_beam_search]
+            )
+
+        with gr.TabItem("vs. Top K Sampling"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("## Inputs ✍️")
+                    gr.Markdown("General options:")
+                    model_id = gr.Text(value="facebook/opt-125m", label="Model Repository")
+                    input_text = gr.Textbox(value="DeepMind Company is", lines=5, label="Input Text")
+                    max_new_tokens = gr.Slider(value=50, minimum=1, maximum=256, label="New tokens to generate")
+                    gr.Markdown("Contrastive Search options:")
+                    alpha = gr.Slider(value=0.6, minimum=0.01, maximum=1.0, step=0.01, label="Alpha")
+                    k = gr.Slider(value=6, minimum=1, maximum=20, step=1, label="K")
+                    gr.Markdown("Sampling options:")
+                    top_k = gr.Slider(value=50, minimum=1, maximum=100, step=1, label="Top K")
+                    seed = gr.Number(value=42, precision=0, label="Seed")
+                    generate_button = gr.Button(value="Generate", label="Generate")
+
+                with gr.Column():
+                    gr.Markdown("## Outputs 🤖")
+                    gr.Markdown("Contrastive Search generation:")
+                    text_contrastive = gr.Textbox(value="", label="")
+                    time_contrastive = gr.Number(value=0.0, precision=1, label="Generation time (ms)")
+                    gr.Markdown("Top K Sampling generation:")
+                    text_top_k = gr.Textbox(value="", label="")
+                    time_top_k = gr.Number(value=0.0, precision=1, label="Generation time (ms)")
+
+            # actions
+            generate_button.click(
+                fn=generate_top_k,
+                inputs=[input_text, model_id, max_new_tokens, alpha, k, top_k, seed],
+                outputs=[text_contrastive, time_contrastive, text_top_k, time_top_k]
+            )
+
+        with gr.TabItem("vs. Nucleus Sampling"):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("## Inputs ✍️")
+                    gr.Markdown("General options:")
+                    model_id = gr.Text(value="facebook/opt-125m", label="Model Repository")
+                    input_text = gr.Textbox(value="DeepMind Company is", lines=5, label="Input Text")
+                    max_new_tokens = gr.Slider(value=50, minimum=1, maximum=256, label="New tokens to generate")
+                    gr.Markdown("Contrastive Search options:")
+                    alpha = gr.Slider(value=0.6, minimum=0.01, maximum=1.0, step=0.01, label="Alpha")
+                    k = gr.Slider(value=6, minimum=1, maximum=20, step=1, label="K")
+                    gr.Markdown("Sampling options:")
+                    top_p = gr.Slider(value=0.95, minimum=0.01, maximum=1.0, step=0.01, label="Top P")
+                    seed = gr.Number(value=42, precision=0, label="Seed")
+                    generate_button = gr.Button(value="Generate", label="Generate")
+
+                with gr.Column():
+                    gr.Markdown("## Outputs 🤖")
+                    gr.Markdown("Contrastive Search generation:")
+                    text_contrastive = gr.Textbox(value="", label="")
+                    time_contrastive = gr.Number(value=0.0, precision=1, label="Generation time (ms)")
+                    gr.Markdown("Nucleus Sampling generation:")
+                    text_top_p = gr.Textbox(value="", label="")
+                    time_top_p = gr.Number(value=0.0, precision=1, label="Generation time (ms)")
+
+            # actions
+            generate_button.click(
+                fn=generate_top_p,
+                inputs=[input_text, model_id, max_new_tokens, alpha, k, top_p, seed],
+                outputs=[text_contrastive, time_contrastive, text_top_p, time_top_p]
+            )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,3 @@
+git+https://github.com/huggingface/transformers@main#egg=transformers
+sentencepiece
+torch