<ADD> +app

app.py

@@ -0,0 +1,174 @@
+import json
+from pathlib import Path
+
+import gradio as gr
+
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+
+from common import setup_cpu
+from models import build_tokenizer, build_model
+from models.meta_optimizer import AttnOptimWrapper
+from tasks import load_task
+from tasks.loader import TokenizedForMCRightPad
+
+DISPLAY_MAPPING = {
+    "sst2": {"positive": "Pos", "negative": "Neg"},
+    "trec": {},
+}
+
+
+@torch.no_grad()
+def do_infer_probs(model, exemplar_attn_kv, exemplar_attn_mask, batched_choices_input):
+    batched_choices_logprobs = []
+    for batched_one_choice_input in batched_choices_input:
+        batch_input_ids, batch_attention_mask, batch_choice_start, batch_choice_end = batched_one_choice_input
+        bs = len(batch_input_ids)
+
+        merged_attn_mask = torch.cat((exemplar_attn_mask.expand(bs, -1), batch_attention_mask), dim=1)
+        # [B, #Heads, Length, Hidden]
+        expand_exemplar_attn_kv = [[layer_k.expand((bs, -1, -1, -1)), layer_v.expand((bs, -1, -1, -1))] for layer_k, layer_v in exemplar_attn_kv]
+
+        batched_logits = model(
+            input_ids=batch_input_ids,  # [B, L']
+            attention_mask=merged_attn_mask,  # [B, L + L']
+            past_key_values=expand_exemplar_attn_kv,  # num_layers * 2 * [B, num_heads, L, H]
+        ).logits
+        batched_output = F.log_softmax(batched_logits, dim=-1)  # [B, L', Vocab]
+
+        batched_one_choice_logprobs = []
+        for input_ids, choice_start, choice_end, lm_logprobs in zip(batch_input_ids, batch_choice_start, batch_choice_end, batched_output):
+            choice_tokens = input_ids[choice_start:choice_end].unsqueeze(1)  # [L, 1]
+            choice_logprobs = lm_logprobs[choice_start - 1 : choice_end - 1]  # [L, Vocab]
+
+            extracted = torch.gather(choice_logprobs, -1, choice_tokens).squeeze(-1)
+
+            choice_length = choice_end - choice_start
+            lm_log_p = torch.sum(extracted).item()
+            norm_lm_log_p = (lm_log_p / choice_length).item()
+
+            choice_lm_info = {"lm_log_p": lm_log_p, "norm_lm_log_p": norm_lm_log_p}
+            batched_one_choice_logprobs.append(choice_lm_info)
+        batched_choices_logprobs.append(batched_one_choice_logprobs)
+    return batched_choices_logprobs
+
+
+@torch.no_grad()
+def process_once(dataset_name, exemplar_str, forward_steps, raw_data):
+    model_name, model_size = "opt", "125m"
+    step_size, momentum = 0.01, 0.9
+
+    setup_cpu(seed=seed)
+    TaskHandler = load_task(dataset_name)
+    task_agent = TaskHandler(prompt_version)
+
+    tokenizer = build_tokenizer(model_name, model_size, padding_side="right")
+    model = build_model(model_name, model_size, False)
+    torch.autograd.set_grad_enabled(False)
+
+    processed_data = task_agent.dataset_preprocess(raw_data)
+    dataset = TokenizedForMCRightPad(processed_data, tokenizer, task_agent.multiple_choice_promptify)
+
+    exemplar_input_ids, exemplar_attn_mask = dataset.tokenize_demonstration(exemplar_str)
+    loader = DataLoader(dataset, shuffle=False, drop_last=False, batch_size=1)
+    meta_optim = AttnOptimWrapper(model, model_name, step_size=step_size, momentum=momentum)
+    meta_optim.init()
+
+    for _ in range(forward_steps):
+        exemplar_kv = meta_optim.step(exemplar_input_ids)
+
+    generated_info = []  # question * [choice0_prob, choice1_prob]
+    for batch_input in loader:
+        batch_output = do_infer_probs(model, exemplar_kv, exemplar_attn_mask.unsqueeze(0), batch_input)  # [batch_of_choice0, batch_of_choice1, ...]
+        zipped_logprobs = list(zip(*batch_output))  # batch * (choice0, choice1, ...)
+        generated_info.extend(zipped_logprobs)
+
+    all_predicted = []
+    for idx, (data, choice_info) in enumerate(zip(processed_data, generated_info)):
+        merged_choice_info = task_agent.merge_choice_info(choice_info)
+        merged_predictions_idx = task_agent.choice_info_to_predictions(merged_choice_info)["lm_log_p"]
+        predicted = task_agent.CHOICES[merged_predictions_idx]
+        ground_truth = task_agent.CHOICES[data["answer_idx"]]
+        res = f"{DISPLAY_MAPPING[dataset_name][predicted]}{'✅' if predicted == ground_truth else '❌'}"
+        all_predicted.append(res)
+    return all_predicted
+
+
+def transpose(l):
+    return list(map(list, zip(*l)))
+
+
+def button_pressed(prev_state):
+    dataset_name = prev_state["dataset_name"]
+    exemplar_str = prev_state["exemplar_str"]
+    forward_steps = prev_state["step"] + 2
+    raw_data = prev_state["raw_data"]
+    prev_table_data = prev_state["table_data"]
+
+    current_output = process_once(dataset_name, exemplar_str, forward_steps, raw_data)
+
+    t_prev = transpose(prev_table_data)
+    t_prev.append([f"T={forward_steps}"] + current_output)
+    updated_table_data = transpose(t_prev)
+
+    ret = [
+        {
+            "dataset_name": dataset_name,
+            "exemplar_str": exemplar_str,
+            "raw_data": raw_data,
+            "step": forward_steps,
+            "table_data": updated_table_data,
+        },
+        f"Step + 2, Now: {forward_steps}",
+        updated_table_data,
+    ]
+    return ret
+
+
+if __name__ == "__main__":
+    dataset_name = "sst2"
+    seed = 0
+    prompt_version = "default"
+    kv_iter = 10
+
+    print(f"Dataset: {dataset_name}")
+    task_root = Path("example_sets").joinpath(dataset_name)
+
+    with task_root.joinpath("demos.txt").open("r") as f:
+        demos = f.read()
+    with task_root.joinpath("sample.pkl").open("r") as f:
+        data = json.load(f)
+        raw_data = [data[str(i)] for i in range(len(data))]
+
+    css = """ #the-table > div > div > div > table > thead {display: none}"""
+
+    title = "🤔 Iterative Forward Tuning Boosts In-context Learning in Language Models"
+    demo = gr.Blocks(css=css, title="🤔Deep-Thinking")
+    with demo:
+        gr.Markdown(f"<h1 style='text-align: center; margin-bottom: 1rem'>{title}</h1>")
+        with gr.Tab("SST-2"):
+            mapping = ["negative", "positive"]
+
+            init_columns = [[e["sentence"], f"*{DISPLAY_MAPPING['sst2'][mapping[e['label']]]}*"] for e in raw_data]
+            state = gr.State(
+                {
+                    "dataset_name": "sst2",
+                    "exemplar_str": demos,
+                    "raw_data": raw_data,
+                    "step": 0,
+                    "table_data": [["**Test Input**", "**Golden**"], *init_columns],
+                }
+            )
+
+            prompt = gr.Textbox(label="Demonstrations (Prompt template formatted)", value=demos)
+            big_table = gr.DataFrame(
+                value=[["**Test Input**", "**Golden**"], *init_columns],
+                elem_id="the-table",
+                datatype=["markdown"] * 50,
+                headers=None,
+            )
+            step_button = gr.Button("Step + 2, Now: 0")
+            step_button.click(button_pressed, inputs=[state], outputs=[state, step_button, big_table])
+
+    demo.launch(server_name="0.0.0.0")

common.py

@@ -0,0 +1,20 @@
+import os
+import random
+
+import numpy as np
+import torch
+
+from tasks import task_mapper
+
+
+def setup_plain_seed(SEED):
+    os.environ["PYTHONHASHSEED"] = str(SEED)
+    random.seed(SEED)
+    np.random.seed(SEED)
+
+
+def setup_cpu(seed):
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    setup_plain_seed(seed)
+    torch.manual_seed(seed)
+    torch.random.manual_seed(seed)

example_sets/sst2/demos.txt

@@ -0,0 +1,6 @@
+Review: and mainly unfunny
+Sentiment: negative
+
+Review: filmmakers david weissman and bill weber benefit enormously from the cockettes ' camera craziness -- not only did they film performances , but
+Sentiment: positive
+

example_sets/sst2/sample.pkl

@@ -0,0 +1,10 @@
+{
+    "0": { "sentence": "the cold turkey would 've been a far better title . ", "label": 0, "idx": 57 },
+    "1": { "sentence": "it 's a cookie-cutter movie , a cut-and-paste job . ", "label": 0, "idx": 28 },
+    "2": { "sentence": "a solid film ... but more conscientious than it is truly stirring . ", "label": 1, "idx": 143 },
+    "3": { "sentence": "it 's slow -- very , very slow . ", "label": 0, "idx": 4 },
+    "4": { "sentence": "filmmakers who can deftly change moods are treasures and even marvels . ", "label": 1, "idx": 679 },
+    "5": { "sentence": "it all adds up to good fun . ", "label": 1, "idx": 393 },
+    "6": { "sentence": "i am sorry that i was unable to get the full brunt of the comedy . ", "label": 0, "idx": 423 },
+    "7": { "sentence": "hilariously inept and ridiculous . ", "label": 1, "idx": 112 }
+}

models/__init__.py

@@ -0,0 +1 @@
+from .huggingface import build_model_signature, build_tokenizer, build_model

models/huggingface.py

@@ -0,0 +1,41 @@
+from transformers import AutoTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
+
+
+def build_model_signature(model_type, model_size):
+    if model_type == "opt":
+        # ["125m", "350m", "1.3b", "2.7b", "6.7b", "13b", "30b", "66b"]
+        return f"facebook/opt-{model_size}"
+    if model_type == "gpt2":
+        # ["sm", "medium", "large", "xl"]
+        if model_size == "sm":
+            return "gpt2"
+        return f"gpt2-{model_size}"
+    if model_type == "e-gpt":
+        # ["neo-125M", "neo-1.3B", "neo-2.7B", "j-6B", "neox-20b"]
+        return f"EleutherAI/gpt-{model_size}"
+    if model_type == "bloom":
+        # ["560m", "1b1", "1b7", "3b", "7b1"]
+        return f"bigscience/bloom-{model_size}"
+
+
+def build_tokenizer(model_type, model_size, padding_side="left", use_fast=False):
+    sign = build_model_signature(model_type, model_size)
+    if not use_fast:
+        tok = AutoTokenizer.from_pretrained(sign, padding_side=padding_side)
+    else:
+        tok = PreTrainedTokenizerFast.from_pretrained(sign, padding_side=padding_side)
+    if model_type in ["gpt2", "e-gpt"]:
+        tok.pad_token_id = tok.eos_token_id
+        tok.pad_token = tok.eos_token
+    return tok
+
+
+def build_model(model_type, model_size, in_8bit):
+    sign = build_model_signature(model_type, model_size)
+    model = AutoModelForCausalLM.from_pretrained(
+        sign,
+        device_map="auto",
+        load_in_8bit=in_8bit,
+    )
+    model.eval()
+    return model

models/meta_optimizer.py

@@ -0,0 +1,78 @@
+import torch
+
+
+class MomentumOptim:
+    def __init__(self, step_size=0.01, momentum=0.9):
+        self.step_size = step_size
+        self.momentum = momentum
+        self.m = None  # velocity
+
+    def init(self):
+        self.m = None
+
+    def upd_m(self, old_m, g):
+        return g + self.momentum * old_m
+
+    def upd(self, old_x, m):
+        return old_x + self.step_size * m
+
+    def __call__(self, old_xs, new_xs):
+        pesudo_gs = [new_x - old_x for old_x, new_x in zip(old_xs, new_xs)]
+
+        if not self.m:
+            self.m = pesudo_gs
+        else:
+            self.m = [self.upd_m(old_m, g) for old_m, g in zip(self.m, pesudo_gs)]
+
+        updated_kv = [self.upd(old_x, m) for old_x, m in zip(old_xs, self.m)]
+        return updated_kv
+
+
+class AttnOptimWrapper:
+    def __init__(self, llm, model_type, optimizer="momentum", **optimizer_args):
+        self.model = llm
+        self.kv = None
+        self.model_type = model_type
+
+        if optimizer == "momentum":
+            self.optim_k = MomentumOptim(**optimizer_args)
+            self.optim_v = MomentumOptim(**optimizer_args)
+        else:
+            raise ValueError()
+
+    def init(self):
+        self.optim_k.init()
+        self.optim_v.init()
+
+    @torch.no_grad()
+    def step(self, ctx_ids):
+        L = len(ctx_ids)
+
+        ctx_ids = ctx_ids.unsqueeze(0)  # [1, L]
+        mask = torch.ones_like(ctx_ids)
+        if self.kv is not None:
+            mask = mask.repeat(1, 2)  # [1, 2*L]
+
+        next_kv = self.model(
+            input_ids=ctx_ids,
+            attention_mask=mask,
+            past_key_values=self.kv,
+            use_cache=True,
+        ).past_key_values  # kv @ (old_ctx + new_ctx)
+
+        cur_kv = []
+        for layer_k, layer_v in next_kv:
+            # [B, num_head, 2*L, head_hidden]
+            cur_kv.append([layer_k[:, :, -L:, :], layer_v[:, :, -L:, :]])  # kv @ (new_ctx)
+
+        if not self.kv:
+            self.kv = cur_kv
+        else:
+            old_ks, old_vs = zip(*self.kv)
+            cur_ks, cur_vs = zip(*cur_kv)
+
+            upd_ks = self.optim_k(old_ks, cur_ks)
+            upd_vs = self.optim_v(old_vs, cur_vs)
+            self.kv = list(zip(upd_ks, upd_vs))
+
+        return self.kv

tasks/__init__.py

@@ -0,0 +1,11 @@
+from .sst2 import SST2ProbInferenceForMC
+
+
+task_mapper = {"sst2": SST2ProbInferenceForMC}
+
+
+def load_task(name):
+    if name not in task_mapper.keys():
+        raise ValueError(f"Unrecognized dataset `{name}`")
+
+    return task_mapper[name]

tasks/base.py

@@ -0,0 +1,58 @@
+import numpy as np
+
+
+class BaseProbInference:
+    def __init__(self, prompt_version):
+        if prompt_version == "default":
+            self.prompt_version = self.default_prompt_version()
+        else:
+            self.prompt_version = prompt_version
+
+        self.raw_data_result = None
+        self.raw_data_sample = None
+        self.raw_data_dev = None
+
+        self.can_be_stratified = False
+        self.CHOICES = None
+        self.num_base_shot = 1
+
+    def default_prompt_version(self):
+        raise NotImplementedError
+
+    def dataset_signature(self):
+        # {
+        #      "result":  (dataset_name, subset, split),  # which produce the final result
+        #      "sample": (dataset_name, subset, split),  # which we sample ICL few-shot examples
+        # }
+        raise NotImplementedError
+
+    def dataset_part(self, part):
+        return self.dataset_signature()[part]
+
+    def dataset_preprocess(self, raw_data):
+        raise NotImplementedError
+
+    def handcrafted_exemplars(self):
+        raise NotImplementedError
+
+    def exemplar_seperator(self):
+        raise NotImplementedError
+
+    def multiple_choice_promptify(self, query, choice):
+        raise NotImplementedError
+
+    @staticmethod
+    def merge_choice_info(choice_info):
+        merged = {}
+        for k in ["lm_log_p", "norm_lm_log_p"]:
+            one_metric_merged = []
+            for info in choice_info:
+                one_metric_merged.append(info[k])
+            merged[k] = one_metric_merged
+        return merged
+
+    @staticmethod
+    def choice_info_to_predictions(info):
+        lm_log_p_idx = int(np.argmax(info["lm_log_p"]))
+        norm_lm_log_p_idx = int(np.argmax(info["norm_lm_log_p"]))
+        return {"lm_log_p": lm_log_p_idx, "norm_lm_log_p": norm_lm_log_p_idx}

tasks/loader.py

@@ -0,0 +1,96 @@
+import torch
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizer
+
+
+class TokenizedForMCRightPad(Dataset):
+    def __init__(self, data, tok: PreTrainedTokenizer, prompt_fn):
+        # data: [query: str, choices: list(str)]
+        self.tok = tok
+        self.prompt_fn = prompt_fn
+        self.max_length = self._find_max_length(data)
+        self.data = self._build_mc_data(data)
+
+    def _find_max_length(self, data):
+        max_len = 0
+
+        def tok_len(t):
+            return len(self.tok.encode(t))
+
+        for ex in data:
+            query = ex["query"]
+            len_choices = [tok_len(self.prompt_fn(query, c)[1]) for c in ex["choices"]]
+            max_len = max(max_len, *len_choices)
+
+        return max_len
+
+    def _build_mc_data(self, data):
+        processed = []
+        num_choices = set(len(e["choices"]) for e in data)
+        if not len(num_choices) == 1:
+            raise ValueError(f"Queries have different number of choices, which is not supported! #choices: {num_choices}")
+        for ex in data:
+            query, choices = ex["query"], ex["choices"]
+            processed_input = [self.prompt_fn(query, choice) for choice in choices]
+            processed_input = [self.tokenize(t_query, t_full) for t_query, t_full in processed_input]
+            processed.append(processed_input)
+
+        return processed
+
+    def tokenize_demonstration(self, demonstration):
+        e = self.tok(demonstration)
+        return torch.LongTensor(e["input_ids"]), torch.LongTensor(e["attention_mask"])  # no padding
+
+    def tokenize(self, only_query, full_text):
+        tok_only_query = self.tok(only_query, add_special_tokens=False)
+        tok_full_no_padding = self.tok(full_text, add_special_tokens=False)
+        tok_full = self.tok(
+            full_text,
+            padding="max_length",
+            max_length=self.max_length,
+            add_special_tokens=False,
+        )  # <pad> is not a special token
+        # tok_only_query = self.tok(only_query)
+        # tok_full_no_padding = self.tok(full_text)
+        # tok_full = self.tok(
+        #     full_text,
+        #     padding="max_length",
+        #     max_length=self.max_length,
+        # )  # <pad> is not a special token
+
+        # print(f"tok_only_query: {self.tok.convert_ids_to_tokens(tok_only_query.input_ids)}")
+        # print(f"tok_full_no_padding: {self.tok.convert_ids_to_tokens(tok_full_no_padding.input_ids)}")
+        # print(f"tok_full: {self.tok.convert_ids_to_tokens(tok_full.input_ids)}")
+        # exit(0)
+
+        len_full = len(tok_full_no_padding.input_ids)
+        len_query = len(tok_only_query.input_ids)
+        e = {
+            "input_ids": tok_full.input_ids,
+            "attention_mask": tok_full.attention_mask,
+            "choice_start": len_query,
+            "choice_end": len_full,
+        }
+        # print("Attn:")
+        # print(tok_full.attention_mask)
+        # print("input_ids:")
+        # print(tok_full.input_ids)
+
+        dcd_sp = self.tok.convert_ids_to_tokens(tok_full.input_ids, skip_special_tokens=False)
+
+        # print(f'{e["choice_start"]}: {e["choice_end"]} = [{self.tok.convert_tokens_to_string(dcd_sp[e["choice_start"] : e["choice_end"]])}]')
+
+        return e
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        def _get_one_item(e):
+            return torch.LongTensor(e["input_ids"]), torch.LongTensor(e["attention_mask"]), e["choice_start"], e["choice_end"]
+
+        es = self.data[idx]
+        # num_choices * (input_ids, attn, start_idx, end_idx)
+        # input_ids, attn: [B, L]
+        # start_idx, end_idx: [B, ]
+        return [_get_one_item(e) for e in es]

tasks/sst2.py

@@ -0,0 +1,43 @@
+from tasks.base import BaseProbInference
+
+
+class SST2ProbInferenceForMC(BaseProbInference):
+    def __init__(self, prompt_version):
+        super().__init__(prompt_version)
+
+        self.CHOICES = ["negative", "positive"]
+        self.can_be_stratified = True
+        self.num_base_shot = len(self.CHOICES)
+
+    def default_prompt_version(self):
+        return "sp"
+
+    def dataset_signature(self):
+        return {
+            "result": ("glue", "sst2", "validation"),
+            "sample": ("glue", "sst2", "train"),
+        }
+
+    def dataset_preprocess(self, raw_data):
+        data = []
+        for e in raw_data:
+            # print(e, flush=True)
+            data.append({"query": e["sentence"].strip(), "choices": self.CHOICES, "answer_idx": e["label"]})
+        return data
+
+    def handcrafted_exemplars(self):
+        raise NotImplementedError
+
+    def exemplar_seperator(self):
+        if self.prompt_version.startswith("sp"):
+            return "\n\n"
+        else:
+            raise ValueError(f"SST2: Not supported prompt_version: {self.prompt_version}")
+
+    def multiple_choice_promptify(self, query, choice):
+        if self.prompt_version.startswith("sp"):
+            with_query = f"Review: {query}\nSentiment:"
+            with_query_and_choice = f"{with_query} {choice}"
+        else:
+            raise ValueError(f"SST2: Not supported prompt_version: {self.prompt_version}")
+        return with_query, with_query_and_choice