Views
No views yet
1pip install gradio
2pip install git+https://github.com/dill-lab/PILS1import gradio as gr
2import torch
3from pils.models import InversionFromHiddenStatesModel
4
5
6MODEL = InversionFromHiddenStatesModel.from_pretrained(
7 "murtaza/pils-32-llama2-chat-7b")
8MODEL.embedder_no_grad=True
9MODEL.embedder.max_new_tokens = 64
10DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11MODEL = MODEL.to(DEVICE)
12
13def invert(user_prompt):
14 global inp
15 sys_prompt = ''
16 inp = MODEL.embedder_tokenizer.apply_chat_template(conversation=[
17 {"role": "system", "content": sys_prompt},
18 {"role": "user", "content": user_prompt},
19 ], add_generation_prompt=True, return_dict=True, return_tensors='pt')
20
21
22 inp = {f"embedder_{k}": v.to(DEVICE) for k, v in inp.items()}
23 output = MODEL.call_embedding_model(**inp)
24 inp['frozen_embeddings'] = output["embeddings"]
25 with torch.inference_mode():
26 out = MODEL.generate(inp, {"max_length": 64})
27 inverted = MODEL.tokenizer.decode(out[0], skip_special_tokens=True)
28 generated = MODEL.embedder_tokenizer.decode(output["chosen_tokens"][0].squeeze(), skip_special_tokens=True)
29 return generated, inverted
30
31
32demo = gr.Interface(
33 fn=invert,
34 inputs=gr.Textbox(label="Secret prompt"),
35 outputs=(gr.Textbox(label="LLM output"), gr.Textbox(label="Inverter guess"))
36 )
37demo.launch(share=True)@misc{nazir2025betterlanguagemodelinversion,
title={Better Language Model Inversion by Compactly Representing Next-Token Distributions},
author={Murtaza Nazir and Matthew Finlayson and John X. Morris and Xiang Ren and Swabha Swayamdipta},
year={2025},
eprint={2506.17090},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2506.17090},
}