adding application

app.py

@@ -242,14 +242,14 @@ def find_last_token_index(full_ids, word_ids):
 @GPU  # this block runs on a job GPU
 def analyse_word(model_name: str, word: str, patchscopes_template: str, context:str = ""):
     try:
-        text = context+ " " + word
+        # text = context+ " " + word
         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(text, add_special_tokens=False)
+        word_token_ids = tokenizer.encode(word, add_special_tokens=False)
 
         # Instantiate Patchscopes retriever
         patch_retriever = PatchscopesRetriever(
@@ -262,7 +262,7 @@ def analyse_word(model_name: str, word: str, patchscopes_template: str, context:
 
         # Run retrieval for the word across all layers (one pass)
         retrieved_words  = patch_retriever.get_hidden_states_and_retrieve_word(
-            text,
+            word,
             num_tokens_to_generate=len(tokenizer.tokenize(word)),
         )[0]
 
@@ -309,14 +309,14 @@ with gr.Blocks(theme="soft") as demo:
         label="Patchscopes prompt (use X as placeholder)",
         value="repeat the following word X twice: 1)X 2)",
     )
-    context_box = gr.Textbox(label="context", value="")
+    # context_box = gr.Textbox(label="context", value="")
     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, word_box, patchscopes_template, context_box],
+        inputs=[model_name, word_box, patchscopes_template], #, context_box],
         outputs=out_html,
     )