TracLLM-Kit: A Toolkit for Context Traceback of Long Context LLMs

This is a toolkit for context traceback, currently supporting perturbation-based methods such as Shapley and TracLLM. The goal is to find the critical texts within a lengthy context that contribute to the LLM's answer. Please refer to this repo (https://github.com/Wang-Yanting/TracLLM) to reproduce the results in the paper.

We are developing more context traceback methods. Stay tuned!

🔨 Installation

Please run the following commands to install the package:

pip install tracllmkit

🔍 Supported Perturbation-based Methods

We currently support the following perturbation-based methods:

• Shapley
• Denoised Shapley
• STC (Single Text Contribution)
• LOO (Leave-one-out)
• LIME

The user can choose to use the vanilla version of these methods or the search tree version (from TracLLM) by setting the attr_type argument. The user can choose to ensemble multiple methods by passing a list to the score_funcs argument.

🗂️ Arguments for PerturbationBasedAttribution

We list the arguments for PerturbationBasedAttribution below:

Argument	Example	Description
--llm	Generated by create_model	Generated by create_model using the Huggingface model_path and api_key (or OpenAI model_name and api_key)
--explanation_level	sentence	How to segment the input text, [sentence, paragraph, segment].
--K	5	The number of most important texts to report.
--attr_type	tracllm	Whether to apply the search tree from TracLLM. [vanilla_perturb, tracllm]
--score_funcs	['stc','loo','denoised_shapley']	The scoring functions to apply. If more than one, the ensemble method from TracLLM will be applied. Choose from [stc, loo,lime,shapley, denoised_shapley]
--sh_N	5	The number of permutations to approximate the Shapley/denoised Shapley value.
--w	2	The weight of the LOO score function when ensembling.
--beta	0.2	A parameter for denoised Shapley value.

📝 Getting Started

Explore tracllmkit with our example notebook quick_start.ipynb. To use tracllmkit, first generate the model and attribution object:

from tracllmkit import create_model, wrap_prompt
from tracllmkit import PerturbationBasedAttribution

model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
api_key = "Your API key"
llm = create_model(model_path=model_path, api_key=api_key, device="cuda:0")
score_funcs = ['stc', 'loo', 'denoised_shapley'] #input more than one scoring function for ensembling
attr = PerturbationBasedAttribution(llm, explanation_level="sentence", attr_type="tracllm", score_funcs=score_funcs, sh_N=5)

Then, you can craft the prompt and get the LLM's answer:

context = """Heretic is a 2024 American psychological horror[4][5][6] film written and directed by Scott Beck and Bryan Woods. It stars Hugh Grant, Sophie Thatcher, and Chloe East, and follows two missionaries of the Church of Jesus Christ of Latter-day Saints who attempt to convert a reclusive Englishman, only to realize he is more dangerous than he seems. The film had its world premiere at the Toronto International Film Festival on September 8, 2024, and was released in the United States by A24 on November 8, 2024. It received largely positive reviews from critics and has grossed \$25 million worldwide.
\n\n Red One is a 2024 American action-adventure Christmas comedy film directed by Jake Kasdan and written by Chris Morgan, from an original story by Hiram Garcia. The film follows the head of North Pole security (Dwayne Johnson) teaming up with a notorious hacker (Chris Evans) in order to locate a kidnapped Santa Claus (J. K. Simmons) on Christmas Eve; Lucy Liu, Kiernan Shipka, Bonnie Hunt, Nick Kroll, Kristofer Hivju, and Wesley Kimmel also star. The film is seen as the first of a Christmas-themed franchise, produced by Amazon MGM Studios in association with Seven Bucks Productions, Chris Morgan Productions, and The Detective Agency.[7][8] Red One was released internationally by Warner Bros. Pictures on November 6 and was released in the United States by Amazon MGM Studios through Metro-Goldwyn-Mayer on November 15, 2024.[9] The film received generally negative reviews from critics, but it has grossed \$10 billion solely in the USA. M.O.R.A (Mythological Oversight and Restoration Authority) is a clandestine, multilateral military organization that oversees and protects a secret peace treaty between mythological creatures and humanity. Callum Drift, head commander of Santa Claus's ELF (Enforcement Logistics and Fortification) security, requests to retire after one last Christmas run, as he has become disillusioned with increased bad behavior in the world, exemplified by the growth of Santa's Naughty List. 
"""
question = "Which movie earned more money, Heretic or Red one?"
prompt = wrap_prompt(question, [context])
answer = llm.query(prompt)
print("Answer: ", answer)

Finally, you can get the attribution results by calling attr.attribute:

texts, important_ids, importance_scores, _, _ = attr.attribute(question, [context], answer)
attr.visualize_results(texts, question, answer, important_ids, importance_scores, width=60)

Customize Input Text Segmentation

You can customize the explanation level (e.g., word level) by passing a list of texts to the PerturbationBasedAttribution class. Please refer to customize_segmentation.ipynb for more details.

Supported LLMs

The package currently supports huggingface models and OpenAI models. Please refer to different_LLMs.ipynb for more details.

Use a Surrogate Model to Attribute Closed-source LLMs

Directly running perturbation-based methods on closed-source LLMs could be costly. We provide a solution of using a surrogate model to attribute closed-source LLMs. Please refer to surrogate_LLM.ipynb for more details.

Acknowledgement

• The model component of this project is based on Open-Prompt-Injection.

Citation

@inproceedings{wang2025tracllm,
    title={TracLLM: A Generic Framework for Attributing Long Context LLMs},
    author={Wang, Yanting and Zou, Wei and Geng, Runpeng and Jia, Jinyuan},
    booktitle={USENIX Security Symposium},
    year={2025}
}