LLMLingua

repository·main·Indexed 27 days ago

A series of prompt compression tools designed to reduce token counts for LLM inference to save costs and accelerate performance. The library includes LLMLingua for general compression, LongLLMLingua for long-context scenarios and RAG, LLMLingua-2 for faster task-agnostic compression via data distillation, and SecurityLingua for detecting malicious intentions in prompts to safeguard LLMs.

Tokens: 8.2K
Snippets: 26
Records: 38
Agent score: 91%

What's inside LLMLingua

Overview of LLMLingua Series
main
The LLMLingua series provides tools to compress prompts using compact, well-trained language models (e.g., GPT2-small, LLaMA-7B). This enables efficient inference with Large Language Models (LLMs) by removing non-essential tokens, achieving up to 20x compression with minimal performance loss.
Key components include:
- LLMLingua: General prompt compression for accelerated inference.
- LongLLMLingua: Specifically designed for long-context scenarios to mitigate the 'lost in the middle' issue and improve RAG performance.
- LLMLingua-2: A faster (3x-6x improvement) task-agnostic compression method using data distillation from GPT-4.
- SecurityLingua: A safety guardrail model that uses security-aware prompt compression to detect jailbreak attacks with negligible overhead.
https://github.com/microsoft/llmlingua/blob/main/README.md
Use phi-2 model in LLMLingua
main
You can use the phi-2 model for prompt compression. Before initializing, you must update your transformers library to the GitHub version.
Setup:
```
pip install -U git+https://github.com/huggingface/transformers.git
```
Usage:
```
llm_lingua = PromptCompressor("microsoft/phi-2")
```
```
llm_lingua = PromptCompressor("microsoft/phi-2")
```
https://github.com/microsoft/llmlingua/blob/main/DOCUMENT.md

Compress prompts with LLMLingua

main

Use the PromptCompressor class to compress prompts by specifying a target token count. You can initialize the compressor with default settings or specify a specific model (e.g., microsoft/phi-2).

To use quantized models like TheBloke/Llama-2-7b-Chat-GPTQ which require <8GB GPU memory, you must first install optimum and auto-gptq via pip.

from llmlingua import PromptCompressor

# Default usage
llm_lingua = PromptCompressor()
compressed_prompt = llm_lingua.compress_prompt(prompt, instruction="", question="", target_token=200)

# Using a specific model
llm_lingua = PromptCompressor("microsoft/phi-2")

# Using a quantized model (requires pip install optimum auto-gptq)
llm_lingua = PromptCompressor("TheBloke/Llama-2-7b-Chat-GPTQ", model_config={"revision": "main"})

Install dependencies for quantized models
main
If you intend to use quantized models (e.g., TheBloke/Llama-2-7b-Chat-GPTQ) which require less than 8GB of GPU memory, you must install optimum and auto-gptq via pip before proceeding.
https://github.com/microsoft/llmlingua/blob/main/README.md
Get help and report issues for LLMLingua
main
To get help with using LLMLingua, refer to the DOCUMENTATION. To report bugs or request new features, use GitHub Issues. It is recommended to search existing issues before filing a new one to prevent duplicates.
https://github.com/microsoft/llmlingua/blob/main/SUPPORT.md

Train a custom LLMLingua-2 compressor

main

To train a compressor for a specific domain, follow these steps to collect, label, and train on distilled data.

1. Data Collection

Format your data as a list of dictionaries containing _idx and _prompt. Use compress.py to instruct GPT-4 to compress the context, then label_word.py to assign labels, and filter.py to remove poor samples.

Commands:

# Compress original context
cd experiments/llmlingua2/data_collection/
python compress.py --load_origin_from <your data path> --chunk_size 512 --compressor llmcomp --model_name gpt-4-32k --save_path <compressed data save path>

# Assign labels
python label_word.py --load_prompt_from <compressed data save path> --window_size 400 --save_path <labeled data save path>

# Filter samples
python filter.py --load_path <labeled data save path> --save_path <kept data save path>

2. Compressor Training

Train a RoBERTa-based compressor using the filtered data.

Command:

cd experiments/llmlingua2/model_training/
python train_roberta.py --data_path <kept data save path>

cd experiments/llmlingua2/model_training/
python train_roberta.py --data_path <kept data save path>

Load pre-collected GPT-4 compression data from Hugging Face
main
You can use the microsoft/MeetingBank-LLMCompressed dataset from Hugging Face to access GPT-4 compression results. The dataset provides two ways to access the data:
1. Concatenated prompts: Access the full original transcript and the final merged compressed result using prompt and compressed_prompt keys.
2. Chunked lists: Access the original chunks and their corresponding compressed chunks using prompt_list and compressed_prompt_list keys.
```
from datasets import load_dataset

# Load the dataset
data = load_dataset("microsoft/MeetingBank-LLMCompressed", split="train")

for idx, sample in enumerate(data):
    # Option 1: Concatenation of all chunks
    prompt = sample["prompt"]
    compressed_prompt = sample["compressed_prompt"]

    # Option 2: Chunk list
    prompt_list = sample["prompt_list"]
    compressed_prompt_list = sample["compressed_prompt_list"]
```
https://github.com/microsoft/llmlingua/blob/main/experiments/llmlingua2/data_collection/README.md

Integrate LLMLingua with LangChain

main

You can integrate (Long)LLMLingua into LangChain using the LLMLinguaCompressor within a ContextualCompressionRetriever. This allows you to compress retrieved documents before passing them to the LLM.

Requirements:

langchain_community
langchain_openai
A compatible compressor model (e.g., openai-community/gpt2).

from langchain.retrievers import ContextualCompressionRetriever
from langchain_community.retrievers.document_compressors import LLMLinguaCompressor
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(temperature=0)

compressor = LLMLinguaCompressor(model_name="openai-community/gpt2", device_map="cpu")
compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor, base_retriever=retriever
)

compressed_docs = compression_retriever.get_relevant_documents(
    "What did the president say about Ketanji Jackson Brown"
)
pretty_print_docs(compressed_docs)

Setup for Data Collection using GPT-4
main
To collect custom data using GPT-4, you need to install openai (version 0.28) and spacy, then download the English language model.
```
pip install openai==0.28

pip install spacy
python -m spacy download en_core_web_sm
```
https://github.com/microsoft/llmlingua/blob/main/experiments/llmlingua2/README.md

Train SecurityLingua on custom data

main

To train a custom SecurityLingua model, follow these steps:

1. Setup Environment

Run the provided setup script:

bash env_setup.sh

2. Build Training Data

First, annotate the prompts using label_word.py, then filter them using filter.py. Note that you can fine-tune the filtering threshold in filter.py to balance performance and security.

python label_word.py \
    --load_prompt_from SecurityLingua/securitylingua-jailbreak-pairs \
    --window_size 400 \
    --save_path ../results/security_lingua/jailbreak_pairs_annotated.pt

python filter.py \
    --load_path ../results/security_lingua/jailbreak_pairs_annotated.pt \
    --save_path ../results/security_lingua/jailbreak_pairs_annotated_filtered.pt

Note: Ensure your dataset follows the format used in SecurityLingua/securitylingua-jailbreak-pairs before parsing.

3. Train the Model

You can perform single-GPU training or multi-GPU training using accelerate.

Single-GPU Training:

python train_roberta.py \
    --data_path ../results/security_lingua/jailbreak_pairs_annotated_filtered.pt \
    --save_path ../results/security_lingua/jailbreak_pairs_annotated_filtered_roberta.pt \
    --model_name microsoft/llmlingua-2-xlm-roberta-large-meetingbank \
    --num_epoch 5 \
    --run_name meetbank_slingua \
    --wandb_project slingua \
    --wandb_name meetbank_slingua

Multi-GPU Training:

ACCELERATE_LOG_LEVEL="ERROR" accelerate launch --num_processes 4 experiments/llmlingua2/model_training/train_roberta.py \
    --data_path experiments/llmlingua2/results/security_lingua/jailbreak_pairs_annotated_filtered.pt  \
    --save_path experiments/llmlingua2/results/models/xlm_slingua.pth \
    --num_epoch 5 \
    --run_name xlm_slingua \
    --wandb_project slingua \
    --wandb_name xlm_slingua

4. Use Custom Checkpoint

Once trained, load your checkpoint into PromptCompressor as shown in the usage guide.

Integrate LLMLingua with LlamaIndex

main

Use LongLLMLinguaPostprocessor in LlamaIndex to compress context during the retrieval process.

Key Parameters for LongLLMLinguaPostprocessor:

instruction_str: The instruction for the final question.
target_token: The target number of tokens.
rank_method: Set to "longllmlingua".
additional_compress_kwargs: A dictionary for advanced settings:
- condition_compare: Boolean.
- condition_in_question: String (e.g., "after").
- context_budget: String (e.g., "+100").
- reorder_context: String (e.g., "sort").
- dynamic_context_compression_ratio: Float.

from llama_index.query_engine import RetrieverQueryEngine
from llama_index.response_synthesizers import CompactAndRefine
from llama_index.indices.postprocessor import LongLLMLinguaPostprocessor

node_postprocessor = LongLLMLinguaPostprocessor(
    instruction_str="Given the context, please answer the final question",
    target_token=300,
    rank_method="longllmlingua",
    additional_compress_kwargs={
        "condition_compare": True,
        "condition_in_question": "after",
        "context_budget": "+100",
        "reorder_context": "sort",  # Enables document reordering
        "dynamic_context_compression_ratio": 0.4, # Enables dynamic compression ratio
    },
)

Collect custom compression data
main
You can construct a custom compression dataset using the provided data collection pipeline script located at data_collection/collect_data.sh.
https://github.com/microsoft/llmlingua/blob/main/experiments/llmlingua2/README.md