adding application

app.py

@@ -1,7 +1,128 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+import pandas as pd
+from functools import lru_cache
 
-def greet(name):
-    return "Hello " + name + "!!"
+# ----------------------------------------------------------------------
+# IMPORTANT: This version uses the PatchscopesRetriever implementation
+# from the Tokens2Words paper (https://github.com/schwartz-lab-NLP/Tokens2Words)
+# ----------------------------------------------------------------------
+try:
+    from word_retriever import PatchscopesRetriever  # pip install tokens2words
+except ImportError:
+    PatchscopesRetriever = None
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+DEFAULT_MODEL = "meta-llama/Llama-3.1-8B"  # light default so the demo boots everywhere
+DEVICE = 'mps'
+#     (
+#     "cuda" if torch.cuda.is_available() else ("mps" if torch.word_retriever.pybackends.mps.is_available() else "cpu")
+# )
+
+@lru_cache(maxsize=4)
+def get_model_and_tokenizer(model_name: str):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.bfloat16 ,
+        output_hidden_states=True,
+    ).to(DEVICE)
+    model.eval()
+    return model, tokenizer
+
+
+def find_last_token_index(full_ids, word_ids):
+    """Locate end position of word_ids inside full_ids (first match)."""
+    for i in range(len(full_ids) - len(word_ids) + 1):
+        if full_ids[i : i + len(word_ids)] == word_ids:
+            return i + len(word_ids) - 1
+    return None
+
+
+def analyse_word(model_name: str, extraction_template: str, word: str, patchscopes_template: str):
+    if PatchscopesRetriever is None:
+        return (
+            "<p style='color:red'>❌ Patchscopes library not found. Run:<br/>"
+            "<code>pip install git+https://github.com/schwartz-lab-NLP/Tokens2Words</code></p>"
+        )
+
+    model, tokenizer = get_model_and_tokenizer(model_name)
+
+    # Build extraction prompt (where hidden states will be collected)
+    extraction_prompt ="X"
+
+    # Identify last token position of the *word* inside the prompt IDs
+    word_token_ids = tokenizer.encode(word, add_special_tokens=False)
+
+    # Instantiate Patchscopes retriever
+    patch_retriever = PatchscopesRetriever(
+        model,
+        tokenizer,
+        extraction_prompt,
+        patchscopes_template,
+        prompt_target_placeholder="X",
+    )
+
+    # Run retrieval for the word across all layers (one pass)
+    retrieved_words  = patch_retriever.get_hidden_states_and_retrieve_word(
+        word,
+        num_tokens_to_generate=len(tokenizer.tokenize(word)),
+    )[0]
+
+    # Build a table summarising which layers match
+    records = []
+    matches = 0
+    for layer_idx, ret_word in enumerate(retrieved_words):
+        match = ret_word.strip(" ") == word.strip(" ")
+        if match:
+            matches += 1
+        records.append({"Layer": layer_idx, "Retrieved": ret_word, "Match?": "✓" if match else ""})
+
+    df = pd.DataFrame(records)
+
+    def _style(row):
+        color = "background-color: lightgreen" if row["Match?"] else ""
+        return [color] * len(row)
+
+    html_table = df.style.apply(_style, axis=1).hide(axis="index").to_html(escape=False)
+
+    sub_tokens = tokenizer.convert_ids_to_tokens(word_token_ids)
+    top = (
+        f"<p><b>Sub‑word tokens:</b> {' , '.join(sub_tokens)}</p>"
+        f"<p><b>Total matched layers:</b> {matches} / {len(retrieved_words)}</p>"
+    )
+    return top + html_table
+
+
+# ----------------------------- GRADIO UI -------------------------------
+with gr.Blocks(theme="soft") as demo:
+    gr.Markdown(
+        """# Tokens→Words Viewer\nInteractively inspect how hidden‑state patching (Patchscopes) reveals a word's detokenised representation across model layers."""
+    )
+
+    with gr.Row():
+        model_name = gr.Dropdown(
+            label="🤖 Model",
+            choices=[DEFAULT_MODEL, "mistralai/Mistral-7B-v0.1", "meta-llama/Llama-2-7b", "Qwen/Qwen2-7B"],
+            value=DEFAULT_MODEL,
+        )
+        extraction_template = gr.Textbox(
+            label="Extraction prompt (use X as placeholder)",
+            value="repeat the following word X twice: 1)X 2)",
+        )
+    patchscopes_template = gr.Textbox(
+        label="Patchscopes prompt (use X as placeholder)",
+        value="repeat the following word X twice: 1)X 2)",
+    )
+    word_box = gr.Textbox(label="Word to test", value="interpretable")
+    run_btn = gr.Button("Analyse")
+    out_html = gr.HTML()
+
+    run_btn.click(
+        analyse_word,
+        inputs=[model_name, extraction_template, word_box, patchscopes_template],
+        outputs=out_html,
+    )
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,213 @@
+accelerate==1.2.1
+aiofiles==23.2.1
+aiohappyeyeballs==2.4.4
+aiohttp==3.11.11
+aiosignal==1.3.2
+annotated-types==0.7.0
+anyio @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_68kdsx8iyd/croot/anyio_1729121281958/work
+appnope @ file:///Users/ktietz/demo/mc3/conda-bld/appnope_1629146036738/work
+argon2-cffi @ file:///opt/conda/conda-bld/argon2-cffi_1645000214183/work
+argon2-cffi-bindings @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_2ef471wnyf/croot/argon2-cffi-bindings_1736182451265/work
+asttokens @ file:///opt/conda/conda-bld/asttokens_1646925590279/work
+async-lru @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_02efro5ps8/croot/async-lru_1699554529181/work
+async-timeout==5.0.1
+attrs @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_93pjmt0git/croot/attrs_1734533120523/work
+Babel @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_00k1rl2pus/croot/babel_1671781944131/work
+backcall @ file:///home/ktietz/src/ci/backcall_1611930011877/work
+beautifulsoup4 @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_94rx5n7wo9/croot/beautifulsoup4-split_1718029832430/work
+bleach @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_faqg19k8gh/croot/bleach_1732292152791/work
+blis==1.2.0
+Brotli @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_f7i0oxypt6/croot/brotli-split_1736182464088/work
+catalogue==2.0.10
+certifi @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_d8j59rqun5/croot/certifi_1734473289913/work/certifi
+cffi @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_e4xd9yd9i2/croot/cffi_1736182819442/work
+charset-normalizer @ file:///croot/charset-normalizer_1721748349566/work
+click==8.1.8
+cloudpathlib==0.20.0
+cloudpickle==3.1.0
+comm @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_3doui0bmzb/croot/comm_1709322861485/work
+confection==0.1.5
+contourpy==1.3.0
+cycler==0.12.1
+cymem==2.0.11
+dask==2024.8.0
+dask-expr==1.1.10
+datasets==3.2.0
+debugpy @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_563_nwtkoc/croot/debugpy_1690905063850/work
+decorator @ file:///opt/conda/conda-bld/decorator_1643638310831/work
+defusedxml @ file:///tmp/build/80754af9/defusedxml_1615228127516/work
+-e git+https://github.com/tokeron/diffusers.git@00769b5d64c2ea35201e0df7a082db3513619afe#egg=diffusers&subdirectory=../../../../../../diffusers
+dill==0.3.8
+distro==1.9.0
+docker-pycreds==0.4.0
+editdistance==0.8.1
+en_core_web_lg @ https://github.com/explosion/spacy-models/releases/download/en_core_web_lg-3.8.0/en_core_web_lg-3.8.0-py3-none-any.whl#sha256=293e9547a655b25499198ab15a525b05b9407a75f10255e405e8c3854329ab63
+en_core_web_md @ https://github.com/explosion/spacy-models/releases/download/en_core_web_md-3.8.0/en_core_web_md-3.8.0-py3-none-any.whl#sha256=5e6329fe3fecedb1d1a02c3ea2172ee0fede6cea6e4aefb6a02d832dba78a310
+en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl#sha256=1932429db727d4bff3deed6b34cfc05df17794f4a52eeb26cf8928f7c1a0fb85
+eval_type_backport==0.2.2
+exceptiongroup @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_b2258scr33/croot/exceptiongroup_1706031391815/work
+executing @ file:///opt/conda/conda-bld/executing_1646925071911/work
+fastapi==0.115.12
+fastjsonschema @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_d1wgyi4enb/croot/python-fastjsonschema_1731939426145/work
+ffmpy==0.5.0
+filelock==3.16.1
+fonttools==4.55.3
+frozenlist==1.5.0
+fsspec==2024.9.0
+gitdb==4.0.12
+GitPython==3.1.44
+gradio==4.44.1
+gradio_client==1.3.0
+h11 @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_110bmw2coo/croot/h11_1706652289620/work
+httpcore @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_fcxiho9nv7/croot/httpcore_1706728465004/work
+httpx @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_cc4egw1482/croot/httpx_1723474826664/work
+huggingface-hub==0.27.1
+idna==3.10
+importlib_metadata @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_cc4qelzghy/croot/importlib_metadata-suite_1732633706960/work
+importlib_resources==6.5.2
+ipykernel @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_ddflobe9t3/croot/ipykernel_1728665605034/work
+ipython @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_6599f73fa7/croot/ipython_1694181355402/work
+jedi @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_38ctoinnl0/croot/jedi_1733987402850/work
+Jinja2 @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_b15nuwux5r/croot/jinja2_1730902833938/work
+jiter==0.8.2
+joblib==1.4.2
+json5 @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_b9ww6ewhv3/croot/json5_1730786813588/work
+jsonschema @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_7boelfqucq/croot/jsonschema_1728486715888/work
+jsonschema-specifications @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_d38pclgu95/croot/jsonschema-specifications_1699032390832/work
+jupyter-events @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_db0avcjzq5/croot/jupyter_events_1718738111427/work
+jupyter-lsp @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_ae9br5v37x/croot/jupyter-lsp-meta_1699978259353/work
+jupyter_client @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_58w2siozyz/croot/jupyter_client_1699455907045/work
+jupyter_core @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_73nomeum4p/croot/jupyter_core_1718818302815/work
+jupyter_server @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_d1t69bk94b/croot/jupyter_server_1718827086930/work
+jupyter_server_terminals @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_e7ryd60iuw/croot/jupyter_server_terminals_1686870731283/work
+jupyterlab @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_a2d0br6r6g/croot/jupyterlab_1725895226942/work
+jupyterlab-pygments @ file:///tmp/build/80754af9/jupyterlab_pygments_1601490720602/work
+jupyterlab_server @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_f64fg3hglz/croot/jupyterlab_server_1725865356410/work
+kiwisolver==1.4.7
+langcodes==3.5.0
+language_data==1.3.0
+locket==1.0.0
+marisa-trie==1.2.1
+markdown-it-py==3.0.0
+MarkupSafe @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_a84ni4pci8/croot/markupsafe_1704206002077/work
+matplotlib==3.9.4
+matplotlib-inline @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_f6fdc0hldi/croots/recipe/matplotlib-inline_1662014472341/work
+matplotlib-venn==1.1.2
+mdurl==0.1.2
+mistune @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_17ya6k1sbs/croots/recipe/mistune_1661496228719/work
+mpmath==1.3.0
+multidict==6.1.0
+multiprocess==0.70.16
+murmurhash==1.0.12
+nbclient @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_626hpwnurm/croot/nbclient_1698934218848/work
+nbconvert @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_f4c1s1qk1f/croot/nbconvert_1728049432295/work
+nbformat @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_2cv_qoc1gw/croot/nbformat_1728049423516/work
+nest-asyncio @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_310vb5e2a0/croot/nest-asyncio_1708532678212/work
+networkx==3.2.1
+nltk==3.9.1
+notebook @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_539v4hufo2/croot/notebook_1727199149603/work
+notebook_shim @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_d6_ze10f45/croot/notebook-shim_1699455897525/work
+numpy==2.0.2
+openai==1.59.7
+orjson==3.10.16
+overrides @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_70s80guh9g/croot/overrides_1699371144462/work
+packaging @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_a6_qk3qyg7/croot/packaging_1734472142254/work
+pandas==2.2.3
+pandocfilters @ file:///opt/conda/conda-bld/pandocfilters_1643405455980/work
+parso @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_8824a1w4md/croot/parso_1733963320105/work
+partd==1.4.2
+patsy==1.0.1
+pexpect @ file:///tmp/build/80754af9/pexpect_1605563209008/work
+pickleshare @ file:///tmp/build/80754af9/pickleshare_1606932040724/work
+pillow==10.4.0
+platformdirs @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_a8u4fy8k9o/croot/platformdirs_1692205661656/work
+plotly==5.24.1
+preshed==3.0.9
+prometheus_client @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_803ymjpv2u/croot/prometheus_client_1731958793251/work
+prompt-toolkit @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_c63v4kqjzr/croot/prompt-toolkit_1704404354115/work
+propcache==0.2.1
+protobuf==5.29.2
+psutil @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_1310b568-21f4-4cb0-b0e3-2f3d31e39728k9coaga5/croots/recipe/psutil_1656431280844/work
+ptyprocess @ file:///tmp/build/80754af9/ptyprocess_1609355006118/work/dist/ptyprocess-0.7.0-py2.py3-none-any.whl
+pure-eval @ file:///opt/conda/conda-bld/pure_eval_1646925070566/work
+pyarrow==18.1.0
+pycparser @ file:///tmp/build/80754af9/pycparser_1636541352034/work
+pydantic==2.10.4
+pydantic_core==2.27.2
+pydub==0.25.1
+Pygments @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_29bs9f_dh9/croot/pygments_1684279974747/work
+pyparsing==3.2.1
+PySocks @ file:///Users/ktietz/Code/oss/ci_pkgs/pysocks_1626781349491/work
+python-box==7.3.0
+python-dateutil @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_66ud1l42_h/croot/python-dateutil_1716495741162/work
+python-json-logger @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_9bjmcmh4nm/croot/python-json-logger_1734370248301/work
+python-multipart==0.0.20
+pytz @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_a4b76c83ik/croot/pytz_1713974318928/work
+PyYAML @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_faoex52hrr/croot/pyyaml_1728657970485/work
+pyzmq @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_95lsut8ymz/croot/pyzmq_1734709560733/work
+referencing @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_5cz64gsx70/croot/referencing_1699012046031/work
+regex==2024.11.6
+requests @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_ee45nsd33z/croot/requests_1730999134038/work
+rfc3339-validator @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_76ae5cu30h/croot/rfc3339-validator_1683077051957/work
+rfc3986-validator @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_d0l5zd97kt/croot/rfc3986-validator_1683058998431/work
+rich==13.9.4
+rpds-py @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_93fzmr7v9h/croot/rpds-py_1732228422522/work
+ruff==0.11.6
+safetensors==0.5.0
+scikit-learn==1.6.0
+scipy==1.13.1
+seaborn==0.13.2
+semantic-version==2.10.0
+Send2Trash @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_5b31f0zzlv/croot/send2trash_1699371144121/work
+sentencepiece==0.2.0
+sentry-sdk==2.19.2
+setproctitle==1.3.4
+shellingham==1.5.4
+six @ file:///tmp/build/80754af9/six_1644875935023/work
+smart-open==7.1.0
+smmap==5.0.2
+sniffio @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_1573pknjrg/croot/sniffio_1705431298885/work
+soupsieve @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_9798xzs_03/croot/soupsieve_1696347567192/work
+spacy==3.8.3
+spacy-legacy==3.0.12
+spacy-loggers==1.0.5
+srsly==2.5.1
+stack-data @ file:///opt/conda/conda-bld/stack_data_1646927590127/work
+starlette==0.46.2
+statsmodels==0.14.4
+swifter==1.4.0
+sympy==1.13.1
+tabulate==0.9.0
+tenacity==9.0.0
+terminado @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_fcfvyc0an2/croot/terminado_1671751835701/work
+thinc==8.3.4
+threadpoolctl==3.5.0
+tinycss2 @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_fcw5_i306t/croot/tinycss2_1668168825117/work
+together==1.4.1
+tokenizers==0.21.0
+tomli @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_d0e5ffbf-5cf1-45be-8693-c5dff8108a2awhthtjlq/croots/recipe/tomli_1657175508477/work
+tomlkit==0.12.0
+toolz==1.0.0
+torch==2.5.1
+torchvision==0.20.1
+tornado @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_0axef5a0m0/croot/tornado_1733960501260/work
+tqdm==4.67.1
+traitlets @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_500m2_1wyk/croot/traitlets_1718227071952/work
+transformers==4.47.1
+typer==0.15.1
+typing_extensions @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_0b3jpv_f79/croot/typing_extensions_1734714864260/work
+tzdata==2024.2
+urllib3 @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_06_m8gdsy6/croot/urllib3_1727769822458/work
+uvicorn==0.34.2
+wandb==0.19.1
+wasabi==1.1.3
+wcwidth @ file:///Users/ktietz/demo/mc3/conda-bld/wcwidth_1629357192024/work
+weasel==0.4.1
+webencodings==0.5.1
+websocket-client @ file:///private/var/folders/nz/j6p8yfhx1mv_0grj5xl4650h0000gp/T/abs_d37u7gqts8/croot/websocket-client_1715878310260/work
+websockets==12.0
+wordcloud==1.9.4
+wrapt==1.17.2
+xxhash==3.5.0
+yarl==1.18.3
+zipp @ file:///private/var/folders/k1/30mswbxs7r1g6zwn8y4fyt500000gp/T/abs_echurpkwug/croot/zipp_1732630743967/work

word_retriever.py

@@ -0,0 +1,189 @@
+import torch
+from tqdm import tqdm
+from abc import ABC, abstractmethod
+
+from .utils.enums import MultiTokenKind, RetrievalTechniques
+from .processor import RetrievalProcessor
+from .utils.logit_lens import ReverseLogitLens
+from .utils.model_utils import extract_token_i_hidden_states
+
+
+class WordRetrieverBase(ABC):
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
+
+    @abstractmethod
+    def retrieve_word(self, hidden_states, layer_idx=None, num_tokens_to_generate=3):
+        pass
+
+
+class PatchscopesRetriever(WordRetrieverBase):
+    def __init__(
+            self,
+            model,
+            tokenizer,
+            representation_prompt: str = "{word}",
+            patchscopes_prompt: str = "Next is the same word twice: 1) {word} 2)",
+            prompt_target_placeholder: str = "{word}",
+            representation_token_idx_to_extract: int = -1,
+            num_tokens_to_generate: int = 10,
+    ):
+        super().__init__(model, tokenizer)
+        self.prompt_input_ids, self.prompt_target_idx = \
+            self._build_prompt_input_ids_template(patchscopes_prompt, prompt_target_placeholder)
+        self._prepare_representation_prompt = \
+            self._build_representation_prompt_func(representation_prompt, prompt_target_placeholder)
+        self.representation_token_idx = representation_token_idx_to_extract
+        self.num_tokens_to_generate = num_tokens_to_generate
+
+    def _build_prompt_input_ids_template(self, prompt, target_placeholder):
+        prompt_input_ids = [self.tokenizer.bos_token_id] if self.tokenizer.bos_token_id is not None else []
+        target_idx = []
+
+        if prompt:
+            assert target_placeholder is not None, \
+                "Trying to set a prompt for Patchscopes without defining the prompt's target placeholder string, e.g., [MASK]"
+
+            prompt_parts = prompt.split(target_placeholder)
+            for part_i, prompt_part in enumerate(prompt_parts):
+                prompt_input_ids += self.tokenizer.encode(prompt_part, add_special_tokens=False)
+                if part_i < len(prompt_parts)-1:
+                    target_idx += [len(prompt_input_ids)]
+                    prompt_input_ids += [0]
+        else:
+            prompt_input_ids += [0]
+            target_idx = [len(prompt_input_ids)]
+
+        prompt_input_ids = torch.tensor(prompt_input_ids, dtype=torch.long)
+        target_idx = torch.tensor(target_idx, dtype=torch.long)
+        return prompt_input_ids, target_idx
+
+    def _build_representation_prompt_func(self, prompt, target_placeholder):
+        return lambda word: prompt.replace(target_placeholder, word)
+
+    def generate_states(self, tokenizer, word='Wakanda', with_prompt=True):
+        prompt = self.generate_prompt() if with_prompt else word
+        input_ids = tokenizer.encode(prompt, return_tensors='pt')
+        return input_ids
+
+    def retrieve_word(self, hidden_states, layer_idx=None, num_tokens_to_generate=None):
+        self.model.eval()
+
+        # insert hidden states into patchscopes prompt
+        if hidden_states.dim() == 1:
+            hidden_states = hidden_states.unsqueeze(0)
+
+        inputs_embeds = self.model.get_input_embeddings()(self.prompt_input_ids.to(self.model.device)).unsqueeze(0)
+        batched_patchscope_inputs = inputs_embeds.repeat(len(hidden_states), 1, 1).to(hidden_states.dtype)
+        batched_patchscope_inputs[:, self.prompt_target_idx] = hidden_states.unsqueeze(1).to(self.model.device)
+
+        attention_mask = (self.prompt_input_ids != self.tokenizer.eos_token_id).long().unsqueeze(0).repeat(
+            len(hidden_states), 1).to(self.model.device)
+
+        num_tokens_to_generate = num_tokens_to_generate if num_tokens_to_generate else self.num_tokens_to_generate
+
+        with torch.no_grad():
+            patchscope_outputs = self.model.generate(
+                do_sample=False, num_beams=1, top_p=1.0, temperature=None,
+                inputs_embeds=batched_patchscope_inputs,# attention_mask=attention_mask,
+                max_new_tokens=num_tokens_to_generate, pad_token_id=self.tokenizer.eos_token_id, )
+
+        decoded_patchscope_outputs = self.tokenizer.batch_decode(patchscope_outputs)
+        return decoded_patchscope_outputs
+
+    def extract_hidden_states(self, word):
+        representation_input = self._prepare_representation_prompt(word)
+
+        last_token_hidden_states = extract_token_i_hidden_states(
+            self.model, self.tokenizer, representation_input, token_idx_to_extract=self.representation_token_idx, return_dict=False, verbose=False)
+
+        return last_token_hidden_states
+
+    def get_hidden_states_and_retrieve_word(self, word, num_tokens_to_generate=None):
+        last_token_hidden_states = self.extract_hidden_states(word)
+        patchscopes_description_by_layers = self.retrieve_word(
+            last_token_hidden_states, num_tokens_to_generate=num_tokens_to_generate)
+
+        return patchscopes_description_by_layers, last_token_hidden_states
+
+
+class ReverseLogitLensRetriever(WordRetrieverBase):
+    def __init__(self, model, tokenizer, device='cuda', dtype=torch.float16):
+        super().__init__(model, tokenizer)
+        self.reverse_logit_lens = ReverseLogitLens.from_model(model).to(device).to(dtype)
+
+    def retrieve_word(self, hidden_states, layer_idx=None, num_tokens_to_generate=3):
+        result = self.reverse_logit_lens(hidden_states, layer_idx)
+        token = self.tokenizer.decode(torch.argmax(result, dim=-1).item())
+        return token
+
+
+class AnalysisWordRetriever:
+    def __init__(self, model, tokenizer, multi_token_kind, num_tokens_to_generate=1, add_context=True,
+                 model_name='LLaMa-2B', device='cuda', dataset=None):
+        self.model = model.to(device)
+        self.tokenizer = tokenizer
+        self.multi_token_kind = multi_token_kind
+        self.num_tokens_to_generate = num_tokens_to_generate
+        self.add_context = add_context
+        self.model_name = model_name
+        self.device = device
+        self.dataset = dataset
+        self.retriever = self._initialize_retriever()
+        self.RetrievalTechniques = (RetrievalTechniques.Patchscopes if self.multi_token_kind == MultiTokenKind.Natural
+                                    else RetrievalTechniques.ReverseLogitLens)
+        self.whitespace_token = 'Ġ' if model_name in ['gemma-2-9b', 'pythia-6.9b', 'LLaMA3-8B', 'Yi-6B'] else '▁'
+        self.processor = RetrievalProcessor(self.model, self.tokenizer, self.multi_token_kind,
+                                            self.num_tokens_to_generate, self.add_context, self.model_name,
+                                            self.whitespace_token)
+
+    def _initialize_retriever(self):
+        if self.multi_token_kind == MultiTokenKind.Natural:
+            return PatchscopesRetriever(self.model, self.tokenizer)
+        else:
+            return ReverseLogitLensRetriever(self.model, self.tokenizer)
+
+    def retrieve_words_in_dataset(self, number_of_examples_to_retrieve=2, max_length=1000):
+        self.model.eval()
+        results = []
+
+        for text in tqdm(self.dataset['train']['text'][:number_of_examples_to_retrieve], self.model_name):
+            tokenized_input = self.tokenizer(text, return_tensors='pt', truncation=True, max_length=max_length).to(
+                self.device)
+            tokens = tokenized_input.input_ids[0]
+            print(f'Processing text: {text}')
+            i = 5
+            while i < len(tokens):
+                if self.multi_token_kind == MultiTokenKind.Natural:
+                    j, word_tokens, word, context, tokenized_combined_text, combined_text, original_word = self.processor.get_next_word(
+                        tokens, i, device=self.device)
+                elif self.multi_token_kind == MultiTokenKind.Typo:
+                    j, word_tokens, word, context, tokenized_combined_text, combined_text, original_word = self.processor.get_next_full_word_typo(
+                        tokens, i, device=self.device)
+                else:
+                    j, word_tokens, word, context, tokenized_combined_text, combined_text, original_word = self.processor.get_next_full_word_separated(
+                        tokens, i, device=self.device)
+
+                if len(word_tokens) > 1:
+                    with torch.no_grad():
+                        outputs = self.model(**tokenized_combined_text, output_hidden_states=True)
+
+                    hidden_states = outputs.hidden_states
+                    for layer_idx, hidden_state in enumerate(hidden_states):
+                        postfix_hidden_state = hidden_states[layer_idx][0, -1, :].unsqueeze(0)
+                        retrieved_word_str = self.retriever.retrieve_word(postfix_hidden_state, layer_idx=layer_idx,
+                                                                          num_tokens_to_generate=len(word_tokens))
+                        results.append({
+                            'text': combined_text,
+                            'original_word': original_word,
+                            'word': word,
+                            'word_tokens': self.tokenizer.convert_ids_to_tokens(word_tokens),
+                            'num_tokens': len(word_tokens),
+                            'layer': layer_idx,
+                            'retrieved_word_str': retrieved_word_str,
+                            'context': "With Context" if self.add_context else "Without Context"
+                        })
+                else:
+                    i = j
+        return results