run_dsl.py

import torch
import json
import sys
from tqdm import tqdm
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers_cfg.token_grammar_recognizer import AbsTokenRecognizer
from backtrack import GenConfig, generate, SequenceState, generation_yielder
from constrainers import Constrainer, TokenRecognizerConstrainer
import gc
from pathlib import Path
from stream_zip import stream_zip, ZIP_64
from datetime import datetime
from stat import S_IFREG
import pickle

MAX_NEW_TOKENS = 512
TEMPERATURE = 1.0
TOP_P = 1.0
TOP_K = 0

@torch.no_grad()
def eval_prob(model, tokenizer, id, input_ids: torch.LongTensor, constrainer: Constrainer, path: Path, method: str, num_iter: int):
    model.eval()

    # Generate sequences
    config = GenConfig(temperature=TEMPERATURE, top_k=TOP_K, top_p=TOP_P)

    def criteria(tokens):
        if tokens[-1] == tokenizer.eos_token_id:
            return SequenceState.FINISHED
        if len(tokens) >= MAX_NEW_TOKENS:
            return SequenceState.INVALID
        return SequenceState.NOT_FINISHED

    pbar = tqdm(range(num_iter), desc=id)
    def set_call_count(count):
        pbar.set_description(f"{id} at {count: 6d}")

    def get_generation():
        return generation_yielder(generate(input_ids, model, constrainer, config), criteria, set_call_count, backtrack=method.startswith('mine'))

    path.mkdir(parents=True, exist_ok=True)
    history = []
    def zip_wrapper():
        timer = -1
        coroutine = get_generation()
        TOKENS, LOGITS = 'tokens', 'logits'
        state = TOKENS
        SUFFIX = {
            TOKENS: 'pkl',
            LOGITS: 'pt',
        }
        while True:
            content = anext(coroutine)
            if isinstance(content, dict):
                history.append({
                    'tokens': content['tokens']
                })
                pbar.update(1)
                content = 'generate'
            yield (
                f'{state}_{timer}.{SUFFIX[state]}',
                datetime.now(),
                S_IFREG | 0o600,
                ZIP_64,
                (pickle.dumps(content),),
            )
            if len(history) >= num_iter:
                break
            if state == LOGITS and method == 'gcd' and not torch.is_tensor(content):
                coroutine = get_generation()
            if state == TOKENS:
                state = LOGITS
            else:
                state = TOKENS
    with open(path / f"{method}_logs.zip", 'xb') as f:
        for chunk in stream_zip(zip_wrapper()):
            f.write(chunk)
    with open(path / f"{method}_samples.jsonl", "x") as f:
        for h in history:
            f.write(json.dumps(h))
            f.write("\n")

    pbar.close()
    constrainer.get_token_acceptance_array_for_stack.cache_clear()

def main():
    sys.setrecursionlimit(10000)
    device = torch.device('cuda')

    MODEL_ID = "mistralai/mistral-7b-instruct-v0.2"

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    tokenizer.pad_token = tokenizer.eos_token

    # Load model
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, device_map=device)
    model.resize_token_embeddings(len(tokenizer))

    datasets = {split: load_dataset("ebmoon/GAD-dataset", split=split) for split in ["SLIA", "BV4", "CP"]}
    datasets['binary'] = [{
        'id': '0',
        'prompt': "Be a helpful assistant. Generate a random binary string of length 5? Directly show the generated string without explanation.",
        'grammar': 'root   ::= s\ns      ::= "10111" | "10001" | "10010" | "11001" | "10011" | "11110" | "10101" | "11111" | "11101" | "00000" | "11100" | "11010" | "10110" | "10000" | "11011" | "10100" | "11000"',
        'samples': 100
    }]

    for split in ["SLIA", "BV4", "CP", 'binary']:
        dataset = datasets[split]
        result_path = Path("results") / split
        for data in dataset:
            id = data['id']
            path = result_path / str(id)
            prompt = data['prompt']
            grammar = data['grammar']
            try:
                samples = data['samples']
            except:
                samples = 2000

            methods = ['mine', 'gad', 'gcd']

            # Tokenize prompt into ids
            input_ids = tokenizer(
                [prompt], add_special_tokens=False, return_tensors="pt", padding=True
            )["input_ids"].squeeze(0)
            input_ids = input_ids.to(model.device)
            # Initialize logits processor for the grammar
            constrainer = TokenRecognizerConstrainer(AbsTokenRecognizer(grammar, tokenizer, "root"), MAX_NEW_TOKENS)
            for method in methods:
                try:
                    eval_prob(model, tokenizer, id, input_ids, constrainer, path, method, samples)
                except FileExistsError:
                    print(f"Skipping {id} {method}")
                gc.collect()

                print(f"Evaluation finished: {id} {method}")


if __name__=="__main__":
    main()