adding application

app.py

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