Hallucination Aware Prompt Compression and Classification

This repository contains work from the Hallucination Prompt Optimization Engine Project, mentored by Archana Vaidheeshwaran and Soham Chatterjee at Summer of Open AI Research, EleutherAI 2025.

Token Classification

The token classification module focuses on evaluating and attributing the importance of different spans of text in a given context to the model's final answer.

Main Logic & Metrics:

• Probs: We measure the last-layer next-token probabilities for each multiple-choice option (A, B, C, D).
• Log-odds: We calculate the log-odds of the probabilities ($\log(p) - \log(1-p)$).
• LOO (Leave-One-Out) Span Attribution: The core logic splits the context into spans (e.g., sentences or paragraphs). It then iteratively removes each span and re-evaluates the log-odds of the target answer. The difference ($\Delta_i$) indicates whether a span is helpful (positive $\Delta$), harmful (negative $\Delta$), or redundant (near zero).

inference.py: This script executes the token classification logic. It loads a language model (e.g., microsoft/Phi-3-mini-128k-instruct or meta-llama/Llama-3.1-8B-Instruct) and evaluates it on the THUDM/LongBench-v2 dataset. It handles prompt building, middle truncation to fit context windows, choice tokenization, greedy inference, and performs the LOO span attribution to output helpful/harmful labels for text spans.

How to Reproduce:

1. Ensure your environment has the required dependencies (see Token_Classification/requirements.txt).
2. Navigate to the Token_Classification directory and run the inference script:

   cd Token_Classification
   python inference.py

Baselines

We tested several baseline prompt compression and optimization models.

Tested Models:

• PCRL (Prompt Compression with Reinforcement Learning): A reinforcement learning approach for compressing prompts.
• RLPrompt: Optimizes discrete text prompts using reinforcement learning.
• SelectiveContext: Evaluates and removes redundant context based on self-information to compress the prompt.
• LLMLingua2: A token-level and context-aware prompt compression method.
• Kaggle Baseline Notebook: Additional baseline implementation.

Baseline Evaluation Pipeline

The repository includes a comprehensive and modular evaluation pipeline (found in Baselines/Baseline- complete pipeline) for testing prompt compression methods like LLMLingua2.

Pipeline Stages:

1. Compression: Uses LLMLingua2 (with support for rate-based, token-based, and context-based methods) to compress the input prompts.
2. Generation: Feeds the compressed prompts to an LLM via the ScaleDown API to generate responses.
3. Evaluation: Assesses the quality of the generated responses against the ground truth using a multi-metric approach.

Dataset:

• The pipeline has built-in integration with the MS-MARCO dataset (v2.1). It can process specific query types such as NUMERIC, BOOLEAN, DATE, and ENTITY.
• It also supports testing on custom datasets provided in JSON format.

Metrics Used:

• LLM-as-a-Judge: Employs an LLM to dynamically evaluate and score the generated response against the ground truth.
• BLEU: Standard string-matching metric measuring n-gram precision.
• ROUGE: Recall-oriented metric for evaluating text generation.

lingua2compressor.ipynb: This Jupyter Notebook provides a quick setup to run the LLMLingua2 compressor standalone. It establishes configurations and connects to the compression API (scaledown.xyz) to compress examples using varying compression methods.

Annotations

Paragraph Level Analysis: The annotations process involves evaluating the context at the paragraph level. Using the LOO span attribution logic, the text is split into paragraphs. Each paragraph is then analyzed and labeled as helpful, harmful, or redundant based on its contribution to the correct model prediction. This analysis helps in understanding which parts of the long context are crucial and which parts might introduce hallucinations or noise.

Contributions

• Chidaksh R - Passage level analysis, Token classification
• Smitali Bhandari- Passage lvl, Baseline pipeline, LLMLingua2, LLMLingua training script
• Jalend Bantupalli- Token lvl classification
• Lakshay Piplani- Token lvl classification
• Shree Bhargavi- PCRL, RLPrompt Baselines
• Eryawan Yulianrifat- Selective Context baseline
• Archana Vaidheeswaran (Mentor)
• Soham Chatterjee (Mentor)

Cite Us

If you use this work in your research, please cite us:

@misc{hallucination_prompt_compression,
  author = {R, Chidaksh and Bhandari, Smitali and Bantupalli, Jalend and Piplani, Lakshay and Bhargavi, Shree and Yulianrifat, Eryawan and Vaidheeswaran, Archana and Chatterjee, Soham},
  title = {Hallucination Aware Prompt Compression and Classification},
  year = {2025},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/ilatims-b/Hallucination-Aware-Prompt-compression}}
}