Add sampling options

app.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Dict, Hashable
+from typing import Any, Dict, Hashable
 
 import streamlit as st
 import streamlit.components.v1 as components
@@ -7,6 +7,7 @@ import numpy as np
 import torch
 import torch.nn.functional as F
 from transformers import AutoModelForCausalLM, AutoTokenizer, BatchEncoding, GPT2LMHeadModel, PreTrainedTokenizer
+from transformers import LogitsProcessorList, RepetitionPenaltyLogitsProcessor, TemperatureLogitsWarper, TopPLogitsWarper, TypicalLogitsWarper
 
 root_dir = Path(__file__).resolve().parent
 highlighted_text_component = components.declare_component(
@@ -118,12 +119,30 @@ window_len = st.select_slider(
 max_tokens = int(MAX_MEM / (multiplier * window_len) - window_len)
 max_tokens = min(max_tokens, 4096)
 
-max_new_tokens = None
+generate_kwargs = {}
 if generation_mode:
-    max_new_tokens = st.slider(
-        "Max. number of generated tokens",
-        min_value=8, max_value=min(1024, max_tokens), value=min(128, max_tokens)
-    )
+    with st.expander("Generation options", expanded=False):
+        generate_kwargs["max_new_tokens"] = st.slider(
+            "Max. number of generated tokens",
+            min_value=8, max_value=min(1024, max_tokens), value=min(128, max_tokens)
+        )
+        col1, col2, col3, col4 = st.columns(4)
+        with col1:
+            generate_kwargs["temperature"] = st.number_input(
+                min_value=0.01, value=0.9, step=0.05, label="`temperature`"
+            )
+        with col2:
+            generate_kwargs["top_p"] = st.number_input(
+                min_value=0., value=0.95, max_value=1., step=0.05, label="`top_p`"
+            )
+        with col3:
+            generate_kwargs["typical_p"] = st.number_input(
+                min_value=0., value=1., max_value=1., step=0.05, label="`typical_p`"
+            )
+        with col4:
+            generate_kwargs["repetition_penalty"] = st.number_input(
+                min_value=1., value=1., step=0.05, label="`repetition_penalty`"
+            )
 
 DEFAULT_TEXT = """
 We present context length probing, a novel explanation technique for causal
@@ -180,13 +199,27 @@ def get_logprobs(model, inputs, metric):
     pbar.empty()
     return logprobs
 
+def get_logits_processor(temperature, top_p, typical_p, repetition_penalty) -> LogitsProcessorList:
+    processor = LogitsProcessorList()
+    if repetition_penalty != 1.0:
+        processor.append(RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty))
+    if temperature != 1.0:
+        processor.append(TemperatureLogitsWarper(temperature))
+    if top_p < 1.0:
+        processor.append(TopPLogitsWarper(top_p=top_p, min_tokens_to_keep=1))
+    if typical_p < 1.0:
+        processor.append(TypicalLogitsWarper(mass=typical_p, min_tokens_to_keep=1))
+    return processor
+
 @torch.inference_mode()
-def generate(model, inputs, metric, window_len, max_new_tokens):
+def generate(model, inputs, metric, window_len, max_new_tokens, **kwargs):
     assert metric == "NLL loss"
     start = max(0, inputs["input_ids"].shape[1] - window_len + 1)
     inputs_window = {k: v[:, start:] for k, v in inputs.items()}
     del inputs_window["labels"]
 
+    logits_warper = get_logits_processor(**kwargs)
+
     new_ids, logprobs = [], []
     eos_idx = None
     pbar = st.progress(0)
@@ -194,9 +227,11 @@ def generate(model, inputs, metric, window_len, max_new_tokens):
     for i in range(max_steps):
         pbar.progress(i / max_steps, f"{i}/{max_steps}")
         inputs_window["attention_mask"] = torch.ones_like(inputs_window["input_ids"], dtype=torch.long)
-        logprobs_window = model(**inputs_window).logits.log_softmax(dim=-1).squeeze(0)
+        logits_window = model(**inputs_window).logits.squeeze(0)
+        logprobs_window = logits_window.log_softmax(dim=-1)
         if eos_idx is None:
-            next_token = torch.multinomial(logprobs_window[-1].exp(), num_samples=1).item()
+            probs_next = logits_warper(inputs_window["input_ids"], logits_window[[-1]]).softmax(dim=-1)
+            next_token = torch.multinomial(probs_next, num_samples=1).item()
             if next_token == tokenizer.eos_token_id or i >= max_new_tokens - 1:
                 eos_idx = i
         else:
@@ -225,7 +260,7 @@ def run_context_length_probing(
     window_len: int,
     metric: str,
     generation_mode: bool,
-    max_new_tokens: int,
+    generate_kwargs: Dict[str, Any],
     cache_key: Hashable
 ):
     del cache_key
@@ -240,7 +275,7 @@ def run_context_length_probing(
                 inputs=_inputs.convert_to_tensors("pt"), 
                 metric=metric,
                 window_len=window_len,
-                max_new_tokens=max_new_tokens
+                **generate_kwargs
             )
             output_ids = [*input_ids, *new_ids]
             window_len = logprobs.shape[1]
@@ -288,7 +323,7 @@ output_ids, scores = run_context_length_probing(
     window_len=window_len,
     metric=metric_name,
     generation_mode=generation_mode,
-    max_new_tokens=max_new_tokens,
+    generate_kwargs=generate_kwargs,
     cache_key=(model_name, text),
 )
 tokens = ids_to_readable_tokens(tokenizer, output_ids)