draftretriever 0.1.1

DraftRetriever

DraftRetriever is an integral component of ADED, a Retrieval-Based Speculative Decoding method that accelerates large language model (LLM) decoding without fine-tuning, using an adaptive draft-verification process. It dynamically adjusts to token probabilities with a tri-gram matrix representation and employs Monte Carlo Tree Search (MCTS) to balance exploration and exploitation, producing accurate drafts quickly. ADED significantly speeds up decoding while maintaining high accuracy, making it ideal for practical applications.

Installation

Prerequisites:

If the provided wheel files are not compatible with your system, ensure you have Rust installed:

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
pip install maturin

Example

Generate Tri-gram Matrix

import draftretriever
from transformers import AutoTokenizer
from tqdm import tqdm
import json

tokenizer = AutoTokenizer.from_pretrained(model_path)


datastore_path = './datastore_chat_large.idx'
writer = draftretriever.Writer(
    file_path=datastore_path,
    vocab_size=tokenizer.vocab_size,
)

dataset_path = "datastore/ShareGPT_V4.3_unfiltered_cleaned_split.json"
assert dataset_path is not None, "please download the dataset from https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered"
dataset = json.load(open(dataset_path))
total_length = len(dataset)
print("number of samples: ", total_length)
for conversations in tqdm(dataset, total=total_length):
    for sample in conversations['conversations']:
        token_list = tokenizer.encode(sample['value'])
        writer.add_entry(token_list)

writer.finalize()

Search

import draftretriever

datastore = draftretriever.Reader(index_file_path=datastore_path)
retrieved_token_list, _draft_attn_mask, _tree_indices, _draft_position_ids, _retrieve_indices = datastore.search(token_list, choices=max_num_draft)

License

Distributed under the MIT License. See LICENSE for more information.

Acknowledgement

The main framework is from REST