init

.gitattributes

@@ -0,0 +1,35 @@
+*.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
+*.ckpt 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
+*.mlmodel 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
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl 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
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* 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

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+.DS_Store
+.vscode

README.md

@@ -0,0 +1,12 @@
+---
+title: Respect
+emoji: 🫡
+colorFrom: pink
+colorTo: yellow
+sdk: gradio
+sdk_version: 4.44.0
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

adapter.py

@@ -0,0 +1,244 @@
+import logging
+import re
+from functools import cache
+from pathlib import Path
+from typing import List, Set, Tuple, TypeVar
+
+import torch
+from PIL import Image
+
+from utils import device, nested_apply, sorted_list
+
+RE_PATTERN = r'^(deselect\s[A-Z](?:\s[A-Z])*(?:\sselect\s[A-Z](?:\s[A-Z])*)?|select\s[A-Z](?:\s[A-Z])*)$'  # noqa
+
+
+# Name type, newtype of str. e.g. "page4-249.png"
+N = TypeVar('N')
+
+ALPHABET = 'ABCDEFGHIJ'  # we only have 10 images
+LEGAL_TOKEN_IDS = [2, 315, 330, 334, 365, 382, 384, 401, 413,
+                   420, 475, 5339, 634, 17960, 32002]  # A - J and <end_of_utterance> and <\s> and 'select' and 'deselect'
+
+
+MINI_DECODER = {
+    384: 'D',
+    # 2: '</s>',
+    32002: '<end_of_utterance>',
+    420: 'G', 17960: 'elect',
+    330: 'A', 365: 'B', 334: 'C', 5339: 'select', 401: 'F', 475: 'J',
+    634: 'des', 315: 'I', 413: 'E', 382: 'H'}
+
+
+class AlphabeticNameHash:
+
+    @cache
+    def __init__(self, context: List[N]) -> None:
+        self._forward_map = {im: ALPHABET[i] for i, im in enumerate(context)}
+        self._backward_map = {ALPHABET[i]: im for i, im in enumerate(context)}
+
+    def hash(self, im: N) -> str:
+        return self._forward_map[im]
+
+    def unhash(self, i: str) -> N:
+        return self._backward_map[i]
+
+    def valid_hash(self, i: str) -> bool:
+        return i in self._backward_map
+
+
+class IdeficsAdapter:
+
+    PAD_TOKEN_ID = 0
+    LABEL_MASK_ID = 32001  # idefics2: image_token_id
+    LEGAL_TOKEN_IDS = LEGAL_TOKEN_IDS
+    LEGAL_TOKEN_MASK = torch.zeros(32003, requires_grad=False)\
+        .index_fill_(0, torch.tensor(LEGAL_TOKEN_IDS), 1).to(device=device(), dtype=torch.bool)
+    SUPPRESS_TOKEN_IDS = list(set(range(32003)) - set(LEGAL_TOKEN_IDS))
+
+    def __init__(self, image_folder: str, processor) -> None:
+        self.t_max_length = 2048
+        self.image_folder = Path(image_folder)
+        self.image_cache = {}
+        self.processor = processor
+        self.tokenizer = self.processor.tokenizer
+
+    def get_image(self, im_name: N) -> Image.Image:
+        if im_name not in self.image_cache:
+            self.image_cache[im_name] = Image.open(
+                self.image_folder.joinpath(im_name))
+        return self.image_cache[im_name]
+
+    def unhash(self, context: List[N], c: str):
+        return AlphabeticNameHash(tuple(context)).unhash(c)
+
+    def valid_hash(self, context: List[N], c: str):
+        return AlphabeticNameHash(tuple(context)).valid_hash(c)
+
+    def parse(self, context: List[N], decoded_out: str,
+              currently_selected: List[N]) -> List[str]:
+        h = AlphabeticNameHash(tuple(context))
+        logging.debug(f"{context=}")
+        # do inference
+        logging.debug(f"{decoded_out=}")
+        selection, deselection = self.parse_raw(decoded_out)
+
+        hashed_currently_selected = {h.hash(n) for n in currently_selected}
+        desel_to_remove = deselection - hashed_currently_selected
+        if len(desel_to_remove) > 0:
+            logging.debug(f"warn! {desel_to_remove=}")
+            deselection = deselection - desel_to_remove
+
+        sel_to_remove = selection & hashed_currently_selected
+        if len(sel_to_remove) > 0:
+            logging.debug(f"warn! {sel_to_remove=}")
+            selection = selection - sel_to_remove
+
+        logging.debug("post strict cleaning")
+        logging.debug(f"{selection=}")
+        logging.debug(f"{deselection=}")
+
+        model_clicks = selection | deselection
+        logging.debug(f"{model_clicks=}")
+        model_clicks_png = [h.unhash(n)
+                            for n in model_clicks if h.valid_hash(n)]
+        logging.debug(f"{model_clicks_png=}")
+        return model_clicks_png
+
+    @staticmethod
+    def parse_raw(text: str) -> Tuple[Set[N], Set[N]]:
+        last_answer = text.strip()
+        if ":" in text:
+            last_answer_pattern = r":.*$"
+            xs = re.findall(last_answer_pattern, text)
+            last_answer = xs[0].removeprefix(":").strip()
+        xs = re.search(RE_PATTERN, last_answer)
+        if xs is None:
+            print(f"{last_answer=}")
+            print("did not pass regex")
+            return set(), set()
+
+        select_pattern = r"(?<!de)select( [A-J])+$"
+        xs = re.search(select_pattern, last_answer)
+        if xs is not None:
+            xs = xs.group()
+        selections = set(xs.split(" ")[1:]) if xs else set()
+
+        deselect_pattern = r"^deselect( [A-J])+"
+        xs = re.search(deselect_pattern, last_answer)
+        if xs is not None:
+            xs = xs.group()
+        deselections = set(xs.split(" ")[1:]) if xs else set()
+
+        return selections, deselections
+
+    def compose(self, context, chats, previous_selected, hash_images, padding):
+        select_accum, deselect_accum, clickss = self.unfold_select_deselect(
+            previous_selected)
+
+        select_accum = select_accum + [[]]
+        deselect_accum = deselect_accum + [[]]
+        previous_selected = [[]] + previous_selected  # old states pre click
+        assert len(chats) == len(select_accum) == len(
+            deselect_accum) == len(previous_selected)
+
+        messages, images = self.build_processor_input(
+            context, chats, select_accum, deselect_accum, previous_selected, hash_images, omit_last_answer=True, sort_names=True, omit_context=False, chat_feedback=None)
+        prompt = self.processor.apply_chat_template(
+            messages, add_generation_prompt=True)
+        prompt = prompt.strip()
+        logging.debug(prompt)
+        # Keep consistent with train_script
+        inputs = self.processor(
+            text=prompt, images=images,
+            padding=padding, truncation=True, max_length=self.t_max_length,
+            return_tensors="pt")
+        return inputs
+
+    def build_processor_input(self, image_pngs: List[N], chats: List[str],
+                              select_accum: List[List[N]],
+                              deselect_accum: List[List[N]],
+                              pre_click_selected_accum: List[List[N]],
+                              hash_image: bool, omit_last_answer: bool,
+                              sort_names: bool, omit_context: bool,
+                              chat_feedback: str, ):
+        def _text_content(text): return {"type": "text", "text": text}
+
+        def _image_content(): return {"type": "image"}
+
+        def _user_prompt(content): return {"role": "user", "content": content}
+
+        def _assistant_prompt(content): return {
+            "role": "assistant", "content": content}
+
+        def _system_prompt(content): return {
+            "role": "system", "content": content}
+
+        def _current_state(selected: List[N]):
+            if len(selected) == 0:
+                return 'none is selected'
+            return f'{" ".join(selected)} currently selected'
+
+        def _listener_action(select: List[N], deselect: List[N]):
+            if len(select) == 0 and len(deselect) == 0:
+                return 'nothing'
+            if len(select) == 0:
+                return f'deselect {" ".join(deselect)}'
+            if len(deselect) == 0:
+                return f'select {" ".join(select)}'
+            return f'deselect {" ".join(deselect)} select {" ".join(select)}'
+
+        func = AlphabeticNameHash(tuple(image_pngs)).hash if hash_image else id
+        context, select_accum, deselect_accum,  pre_click_selected_accum = nested_apply(
+            func, (image_pngs, select_accum, deselect_accum, pre_click_selected_accum))
+
+        prompt = []
+        images = []
+        if not omit_context:
+            images = [self.get_image(im) for im in image_pngs]
+            images_and_names_content = []
+            for im_name in context:
+                images_and_names_content.append(_image_content())
+                images_and_names_content.append(_text_content(im_name))
+            prompt.append(_system_prompt(images_and_names_content))
+        if not len(chats) == len(select_accum) == len(deselect_accum) == len(pre_click_selected_accum):
+            logging.error(f"{chats=}")
+            logging.error(f"{select_accum=}")
+            logging.error(f"{deselect_accum=}")
+            logging.error(f"{pre_click_selected_accum=}")
+            assert False
+        for i, (chat, select, deselect, pre_click_selected) in enumerate(
+            zip(chats, select_accum, deselect_accum, pre_click_selected_accum)):
+            if sort_names:
+                select = sorted(select)
+                deselect = sorted(deselect)
+                pre_click_selected = sorted(pre_click_selected)
+
+            prompt.append(_system_prompt(
+                [_text_content(_current_state(pre_click_selected))]))
+            prompt.append(_user_prompt([_text_content(chat)]))
+            prompt.append(_assistant_prompt(
+                [_text_content(_listener_action(select, deselect))]))
+        if omit_last_answer:
+            # idefics2 has processor.apply_chat_template(messages, add_generation_prompt=True) instead
+            prompt.pop(-1)
+        if chat_feedback is not None:
+            prompt.append(_user_prompt([_text_content(chat_feedback)]))
+        return prompt, images
+
+    def unfold_select_deselect(self, previous_selected: List[List[N]]) -> Tuple[List[N], List[N], List[N]]:
+        # currently selected AFTER i-th turn
+        num_turns = len(previous_selected)
+        selected: List[List[str]] = []  # turn-wise selection
+        deselected: List[List[str]] = []  # turn-wise deselection
+        clicks: List[List[str]] = []
+        # combining turn-wise newly selected and newly deselected
+        prev_selected = set()
+        for turn in range(num_turns):
+            curr_selected = set(previous_selected[turn])
+            newly_selected = curr_selected - prev_selected
+            newly_deselected = prev_selected - curr_selected
+            selected.append(sorted_list(newly_selected))
+            deselected.append(sorted_list(newly_deselected))
+            clicks.append(sorted_list(newly_selected | newly_deselected))
+            prev_selected = curr_selected.copy()
+        return selected, deselected, clicks

app.py

@@ -0,0 +1,335 @@
+import dataclasses
+import logging
+import os
+from typing import Any, Dict, List
+
+import gradio as gr  # type: ignore
+import PIL.Image as Image
+import PIL.ImageOps as ImageOps
+import spaces  # type: ignore
+import torch
+from peft import PeftModel  # type: ignore
+from transformers import AutoProcessor  # type: ignore
+from transformers import Idefics2ForConditionalGeneration, Idefics2Processor
+
+from adapter import IdeficsAdapter
+from config_generator import GameConfig, generate_game_config
+from utils import device, nested_to_device, sorted_list
+import copy
+
+### Constants
+css="""
+.radio-group .wrap {
+    display: grid;
+    grid-template-columns: repeat(5, 1fr);
+    grid-template-rows: repeat(5, 1fr);
+    width: 100%;
+    height: 100%
+}
+"""
+IMG_DIR = "tangram_pngs"
+
+
+### Bot server
+
+GEN_KWS: Dict[str, Any] = {
+            "max_new_tokens": 10,
+            "do_sample": True,
+            "temperature": 1.0,
+            "output_logits": True,
+            "return_dict_in_generate": True,
+            "remove_invalid_values": True,  # just to be safe
+            "renormalize_logits": True,
+            "suppress_tokens": IdeficsAdapter.SUPPRESS_TOKEN_IDS
+        }
+
+@spaces.GPU(duration=20)
+def get_model_response(  # predict
+    model: PeftModel, adapter_name: str, adapter: IdeficsAdapter,
+    image_paths: List[str], chat : str, chats: List[str],
+    previous_selected: List[List[str]]
+) -> List[str]:
+    if model.active_adapter != adapter_name:
+        model.set_adapter(adapter_name)
+
+    model.to(device())
+
+    new_chats = chats + [chat]
+    currently_selected = previous_selected[-1] if len(previous_selected) > 0 else []
+    model_input: Dict[str, Any] = adapter.compose(  # type: ignore
+        image_paths, new_chats, previous_selected, True, False)
+    model_input = nested_to_device(model_input)  # type: ignore
+
+    with torch.inference_mode(), torch.autocast(device_type=device().type,
+                                                dtype=torch.bfloat16):
+        model_output = model.generate(**model_input, **GEN_KWS)  # type: ignore
+
+    decoded_out: str = adapter.tokenizer.decode(  # type: ignore
+        model_output.sequences[0], skip_special_tokens=True)
+    model_clicks = adapter.parse(
+        image_paths, decoded_out, currently_selected)  # type: ignore
+
+    if len(model_clicks) == 0:
+        logging.warning("empty clicks by model")
+        model_clicks = [image_paths[0]]
+        logging.debug(f"{image_paths=}")
+        logging.debug(f"selecting {model_clicks}")
+        prob = -1
+    else:
+        prob = -3
+        logging.debug(f"{prob=}")
+    logging.info(f"User input: {chat}")
+    logging.info(f"Model selected: {model_clicks}")
+    logging.debug(f"Model output: {decoded_out}")
+    return model_clicks
+
+
+def get_model() -> PeftModel:
+    model_id = 'lil-lab/respect'
+    checkpoint = "HuggingFaceM4/idefics2-8b"
+    model = Idefics2ForConditionalGeneration.from_pretrained(  # type: ignore
+        checkpoint, torch_dtype=torch.bfloat16,
+    )
+    peft_model = PeftModel.from_pretrained(  # type: ignore
+        model, model_id, adapter_name="r6_bp", is_trainable=False, revision="r6_bp")
+
+    # Add other adapter - hack to avoid conflict
+    lora_config = copy.deepcopy(peft_model.active_peft_config)
+    targets = list(set(n[:n.find('lora')-1] for n, _ in model.named_parameters()
+                       if 'lora' in n))
+    lora_config.target_modules = targets
+    peft_model.add_adapter("r0", lora_config)
+    peft_model.load_adapter(model_id, "r0", is_trainable=False, revision="r0",
+                            peft_config=lora_config)
+    return peft_model
+
+def get_processor() -> Idefics2Processor:
+    checkpoint = "HuggingFaceM4/idefics2-8b"
+    processor = AutoProcessor.from_pretrained(  # type: ignore
+        checkpoint, do_image_splitting=False,
+        size={"longest_edge": 224, "shortest_edge": 224})
+    return processor # type: ignore
+
+def get_adapter() -> IdeficsAdapter:
+    processor = get_processor()
+    return IdeficsAdapter(IMG_DIR, processor)
+
+
+### Game logic
+
+@dataclasses.dataclass(frozen=False)
+class GameState:
+    config: GameConfig
+    adapter_name: str
+    chats: List[str]
+    currently_selected: List[str]
+    selected_accum: List[List[str]]
+    clicks_accum: List[List[str]]
+    turn: int = 0
+
+    def has_ended(self):
+        return self.has_successfully_ended() or self.turn >= 10
+
+    def has_successfully_ended(self):
+        return set(self.currently_selected) == set(self.config.targets)
+
+    ### UI helpers
+
+    def serialize_conversation(self):
+        output = [f"Turn {i+1}: {message}"
+                  for i, message in enumerate(self.chats)]
+        return "\n".join(output)
+
+    def markup_images(self):
+        context = self.config.speaker_context
+        targets = self.config.targets
+        selected = self.currently_selected
+        changes = self.selected_accum[-1] if len(self.selected_accum) > 0 else []
+
+        tangram_list = self._display_context(context, targets, changes, selected)
+        # return [(img, f"Image {i+1}") for i, img in enumerate(tangram_list)]
+        return tangram_list
+
+    @staticmethod
+    def _display_context(context: List[str], targets: List[str],
+                        changes: List[str], selected: List[str]) -> List[Image.Image]:
+        tangram_list: List[Image.Image] = []
+        arrow = Image.open("yellow_circle.png").resize((20, 20)).convert("RGBA")
+        for img in context:
+            image = Image.open(os.path.join(IMG_DIR, img)).resize((60, 60)).convert("RGB")
+            image = ImageOps.expand(image, border=2, fill="white")
+            if img in targets and img in selected:  # listener selected a target image
+                image = ImageOps.expand(image, border=10, fill="green")
+            elif img in targets and img not in selected:  # unselected target:
+                image = ImageOps.expand(image, border=10, fill="black")
+            elif img in selected and img not in targets:  # listener selected a wrong image
+                image = ImageOps.expand(image, border=10, fill="red")
+            else:
+                image = ImageOps.expand(image, border=10, fill="white")
+            image = ImageOps.expand(image, border=2, fill="white")
+            if img in changes:
+                image.paste(arrow, (68, 0), mask=arrow)
+            tangram_list.append(image)
+        return tangram_list
+
+
+class GameFlow:
+
+    @classmethod
+    def initialize(cls, model_iteration: str) -> GameState:
+        config = generate_game_config()
+        adapter_name = "r0" if model_iteration == "Initial System" else "r6_bp"
+        state = GameState(
+            config=config,
+            adapter_name=adapter_name,
+            chats=[],
+            currently_selected=[],
+            selected_accum=[],
+            clicks_accum=[],
+            turn=0,
+        )
+        return state
+
+    @classmethod
+    def progress(cls, state: GameState, chat: str,
+                      model: PeftModel,
+                      adapter: IdeficsAdapter) -> GameState:
+        turn = state.turn
+        model_context_images = state.config.listener_context
+
+        model_clicks = get_model_response(
+            model, state.adapter_name, adapter,
+            model_context_images, chat,
+            state.chats, state.selected_accum
+        )
+
+        # symmetric difference (apply deselection, then selection)
+        currently_selected2 = sorted_list(
+            (set(state.currently_selected) - set(model_clicks)) \
+            | (set(model_clicks) - set(state.currently_selected))
+        )
+
+        state2 = GameState(
+            # constants
+            config=state.config,
+            adapter_name=state.adapter_name,
+            # updates
+            chats=state.chats.copy() + [chat],
+            currently_selected=currently_selected2,
+            selected_accum=state.selected_accum.copy() + [currently_selected2],
+            clicks_accum=state.clicks_accum.copy() + [model_clicks],
+            turn=turn+1,
+        )
+        return state2
+
+
+
+### UI
+
+def create_app_inner():
+    ### layout
+    gr.Markdown("# Tangram Multi-Reference Game")
+    gr.Markdown(
+        '### You will be playing a multi-reference games against a model. \
+        To start a game, first select whether you wish to play against our \
+        initial trained model ("Initial System") or \
+        our model at the end of continual learning ("Final System") \
+        and press the "Start Game" button. \
+        You will take on a "speaker" role at each round. \
+        Your goal is to describe this image (via a message in the textbox) \
+        so that the model can guess what it is.'
+    )
+
+    gr.Markdown("Targets have black borders. Correctly selected targets have green borders. Incorrectly selected targets have red borders. Actions are marked with yellow dot.")
+
+    gr.Markdown("The listener cannot see boxes or colors and the order is different.")
+
+    gr.Markdown(
+        '### Press "Send" to submit your action to proceed to the next turn. \
+        You have 10 turns in total.'
+    )
+
+    with gr.Row():
+        model_iteration = gr.Radio(["Initial System", "Final System"],
+                                   label="Model Iteration",
+                                   value="Final System")
+        start_btn = gr.Button("Start Game")
+
+    with gr.Row():
+        current_turn = gr.Textbox(label="TURN")
+        success = gr.Textbox(label="Success")
+
+    with gr.Row():
+        image_output = gr.Gallery(
+            label="CONTEXT", show_label=False, elem_id="gallery",
+            columns=5, rows=2, object_fit="contain", height="250px",
+            allow_preview=False, container=True, interactive=False
+        )
+
+    with gr.Row():
+        conversation_output = gr.Textbox(label="Interaction History")
+        user_input = gr.Textbox(label="Your Message as Speaker", interactive=True)
+
+    send_btn = gr.Button("Send", interactive=True)
+
+    ### globals
+    model = get_model()
+    adapter = get_adapter()
+    game_state = gr.State(value=None)
+
+    ### callbacks
+    def output_from_state(state: GameState):
+        has_ended = state.has_ended()
+        success = "success" if state.has_successfully_ended() else "failure"
+        return (
+            state.markup_images(),  # image_output
+            state.serialize_conversation(),  # conversation_output
+            f"{state.turn+1}/10",  # current_turn
+            success if has_ended else "n/a",  # success
+            gr.update(interactive=not has_ended, value=""),  # user_input
+            gr.update(interactive=not has_ended),  # send_btn
+            gr.update(interactive=has_ended),  # model_iteration
+            state,  # game_history
+        )
+
+    def on_start_interaction(model_iteration: str):
+        assert model_iteration in ["Initial System", "Final System"]
+        state = GameFlow.initialize(model_iteration)
+        return output_from_state(state)
+
+    def on_send_message(message: str, state: GameState):
+        nonlocal model
+        nonlocal adapter
+        if message.strip() == "":
+            logging.info("Empty message")
+            return output_from_state(state)
+        state = GameFlow.progress(state, message, model, adapter)
+        return output_from_state(state)
+
+    start_btn.click(
+        on_start_interaction,
+        inputs=[model_iteration],
+        outputs=[image_output, conversation_output, current_turn, success,
+            user_input, send_btn, model_iteration, game_state],
+        queue=False
+    )
+
+    send_btn.click(
+        on_send_message,
+        inputs=[user_input, game_state],
+        outputs=[image_output, conversation_output, current_turn, success,
+                 user_input, send_btn, model_iteration, game_state],
+        queue=True
+    )
+
+
+def create_app():
+    with gr.Blocks(css=css) as app:
+        create_app_inner()
+    return app
+
+
+if __name__ == "__main__":
+    app = create_app()
+    app.queue()
+    app.launch()

config_generator.py

@@ -0,0 +1,51 @@
+import dataclasses
+import functools
+import logging
+import os
+import pickle
+import pprint
+import random
+from typing import List
+
+EMPTY_DATA_PATH =  "tangram_pngs/"
+SPLIT_PATH = "dataset_splits/"
+
+
+@dataclasses.dataclass(frozen=True)
+class GameConfig:
+    speaker_context: List[str]
+    listener_context: List[str]
+    targets: List[str]
+
+
+def generate_game_config() -> GameConfig:
+    corpus = _get_data()
+    context = random.sample(corpus, 10)
+    num_targets = random.randint(3, 5)
+    targets = random.sample(context, num_targets)
+    listener_order = list(range(10))
+    random.shuffle(listener_order)
+
+    config = GameConfig(
+        speaker_context=context,
+        listener_context=[context[i] for i in listener_order],
+        targets=targets,
+    )
+    logging.info(f"context_dict: {pprint.pformat(dataclasses.asdict(config))}")
+    return config
+
+@functools.cache
+def _get_data(restricted_dataset: bool=False):
+    if not restricted_dataset:
+        # 1013 images
+        paths = os.listdir(EMPTY_DATA_PATH)
+    else:
+        # 912 images
+        with open(os.path.join(SPLIT_PATH, "test_imgs.pkl"), 'rb') as f:
+            paths = pickle.load(f)
+        with open(os.path.join(SPLIT_PATH, "train_imgs.pkl"), 'rb') as f:
+            paths += pickle.load(f)
+        paths = [path + ".png" for path in paths]
+    dup_images = ["page6-51.png", "page6-66.png", "page4-170.png"]
+    paths = [path for path in paths if path != ".DS_Store" and path not in dup_images]
+    return paths

requirements.txt

@@ -0,0 +1,9 @@
+torch==2.2.0
+datasets==2.18.0
+transformers==4.40.0
+accelerate==0.29.2
+loralib==0.1.2
+peft==0.10.0
+nltk==3.8.1
+gradio==4.44.1
+spaces==0.30.4