Evaluation: Add more evals (#283)

awq/evaluation/__init__.py

+from awq.evaluation.eval_utils import (
+    evaluate_perplexity,
+    eval_librispeech,
+    eval_mmlu,
+)
+from awq.evaluation.humaneval_utils import eval_humaneval
+from awq.evaluation.kl_divergence import eval_kl_divergence
\ No newline at end of file

awq/evaluation/eval_utils.py

+import torch
+import torch.nn as nn
+from tqdm import tqdm
+from lm_eval import evaluator
+from datasets import load_dataset
+from transformers import pipeline
+from evaluate import load as load_metric
+from lm_eval.tasks import initialize_tasks
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.models.whisper.english_normalizer import BasicTextNormalizer
+
+def get_device():
+    if torch.backends.mps.is_available():
+        return 'mps'
+    elif torch.cuda.is_available():
+        return 'cuda:0'
+    else:
+        return 'cpu'
+
+def evaluate_perplexity(model, tokenizer):
+    def _perplexity(nlls, n_samples, seqlen):
+        return torch.exp(torch.stack(nlls).sum() / (n_samples * seqlen))
+    
+    # load and prepare dataset
+    data = load_dataset('wikitext', 'wikitext-2-raw-v1', split='test')
+    data = tokenizer("\n\n".join(data['text']), return_tensors='pt')
+    data = data.input_ids.to(model.device)
+
+    seqlen = 2048
+    model = model.eval()
+    n_samples = data.numel() // seqlen
+    
+    nlls = []
+
+    with tqdm(range(n_samples), desc="Perplexity -") as progress_bar:
+        for i in progress_bar:
+            start_index = (i * seqlen)
+            end_index = ((i + 1) * seqlen)
+            batch = data[:, start_index:end_index].to(model.device)
+            with torch.no_grad():
+                logits = model(batch).logits
+            shift_logits = logits[:, :-1, :].contiguous().float()
+            shift_labels = data[:, start_index:end_index][:, 1:]
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+            neg_log_likelihood = loss.float() * seqlen
+            nlls.append(neg_log_likelihood)
+
+            curr_ppl = _perplexity(nlls, i+1, seqlen)
+            progress_bar.set_description(f"Perplexity {curr_ppl:.3f}")
+
+    ppl = _perplexity(nlls, n_samples, seqlen)
+    
+    return ppl.item()
+
+def eval_librispeech(model_id, num_samples=100, batch_size=4):
+    try:
+        import jiwer, librosa, soundfile
+    except ImportError:
+        print("Please install the following: pip install jiwer librosa soundfile")
+    
+    dataset = load_dataset("librispeech_asr", "clean", split="test", streaming=True)
+
+    # Load the Whisper model pipeline for automatic speech recognition
+    pipe = pipeline(
+        task="automatic-speech-recognition",
+        model=model_id,
+        batch_size=batch_size,
+        device=get_device(),
+        torch_dtype=torch.float16,
+    )
+
+    # Word normalizer
+    normalizer = BasicTextNormalizer()
+    
+    # Load the WER metric
+    wer_metric = load_metric("wer")
+
+    texts = []
+    audio = []
+    for i, data in tqdm(enumerate(dataset), total=num_samples, desc="Loading dataset"):
+        if len(audio) == num_samples: break
+        audio.append(data["audio"])
+        texts.append(data["text"])
+
+    references = []
+    predictions = []
+
+    with tqdm(range(0, num_samples, batch_size), desc="Word Error Rate: -") as pbar:
+        for i in pbar:
+            batch_audio = audio[i:i+batch_size]
+            batch_texts = texts[i:i+batch_size]
+
+            # inference
+            results = pipe(batch_audio, batch_size=len(batch_audio))
+
+            # normalize text
+            normalized_predictions = [normalizer(result["text"]) for result in results]
+            normalized_texts = [normalizer(text) for text in batch_texts]
+
+            predictions.extend(normalized_predictions)
+            references.extend(normalized_texts)
+
+            # word error rate computation
+            wer = wer_metric.compute(predictions=predictions, references=references) * 100
+            pbar.set_description(f"Word Error Rate: {wer:.3f}%")
+
+def eval_mmlu(model_path="gpt2", num_fewshot=1, batch_size=1, device="cuda:0", task_use_pretrained=False):
+    try:
+        import vllm
+        VLLM_INSTALLED = True
+    except ImportError:
+        VLLM_INSTALLED = False
+    
+    initialize_tasks(verbosity="DEBUG")
+
+    if VLLM_INSTALLED:
+        model = "vllm"
+        model_args = dict(
+            pretrained=model_path,
+            max_model_len=2048,
+            dtype="float16",
+            trust_remote_code=True,
+        )
+
+        if not task_use_pretrained:
+            model_args["quantization"] = "awq"
+    else:
+        model = "hf"
+        model_args = dict(
+            pretrained=model_path,
+            device_map_option=device,
+            dtype="float16",
+            trust_remote_code=True,
+        )
+    model_args = ",".join([f"{k}={v}" for k,v in model_args.items()])
+
+    results = evaluator.simple_evaluate(
+        model=model,
+        model_args=model_args,
+        tasks=['mmlu'],
+        num_fewshot=num_fewshot,
+        batch_size=batch_size,
+        device=device,
+        log_samples=False,
+    )
+
+    print(evaluator.make_table(results))
+
+if __name__ == '__main__':
+    ### PERPLEXITY
+    # model_path = 'mistralai/Mistral-7B-Instruct-v0.1'
+    # model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
+    # tokenizer = AutoTokenizer.from_pretrained(model_path)
+    # evaluate_perplexity(model, tokenizer)
+
+    ### WORD ERROR RATE
+    # model_id = "distil-whisper/distil-small.en" # 3.594
+    model_id = "distil-whisper/distil-medium.en" # 3.436
+    eval_librispeech(model_id)

awq/evaluation/humaneval_utils.py

+# Adapted from https://github.com/abacaj/code-eval
+
+import io
+import os
+import json
+import gzip
+import torch
+import signal
+import tempfile
+import itertools
+import contextlib
+import numpy as np
+from tqdm import tqdm
+import multiprocessing
+from datasets import load_dataset
+from collections import defaultdict, Counter
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+from typing import (
+    List,
+    Union,
+    Iterable,
+    Dict,
+    Optional,
+)
+
+from transformers import (
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+
+def eval_humaneval(
+    model: PreTrainedModel,
+    tokenizer: PreTrainedTokenizer,
+    out_path: str = "humaneval_out.jsonl",
+    format_tabs: bool = True,
+):
+    problems = {example["task_id"]: example for example in load_dataset("openai_humaneval")["test"]}
+    
+    samples = []
+
+    for i, (task_id, task) in tqdm(enumerate(problems.items()), total=len(problems) ):
+        if format_tabs:
+            prompt = task["prompt"].replace("    ", "\t")
+        else:
+            prompt = task["prompt"]
+
+        batch_completions = generate_batch_completion(
+            model, tokenizer, prompt, 1
+        )
+
+        for sample in batch_completions:
+            result = dict(
+                task_id=task_id,
+                completion=sample,
+            )
+
+            samples += [result]
+
+    with open(out_path, 'wb') as fp:
+        for x in samples:
+            fp.write((json.dumps(x) + "\n").encode('utf-8'))
+    
+    results = evaluate_functional_correctness(
+        sample_file=out_path,
+        k=[1],
+        n_workers=4,
+        timeout=3.0,
+    )
+
+    print(results)
+
+@torch.inference_mode()
+def generate_batch_completion(
+    model: PreTrainedModel, tokenizer: PreTrainedTokenizer, prompt, batch_size
+) -> list[str]:
+    input_batch = [prompt for _ in range(batch_size)]
+    inputs = tokenizer(input_batch, return_tensors="pt").to(model.model.device)
+    input_ids_cutoff = inputs.input_ids.size(dim=1)
+
+    generated_ids = model.generate(
+        **inputs,
+        use_cache=True,
+        max_new_tokens=512,
+        temperature=0.2,
+        top_p=0.95,
+        do_sample=True,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+
+    batch_completions = tokenizer.batch_decode(
+        [ids[input_ids_cutoff:] for ids in generated_ids],
+        skip_special_tokens=True,
+    )
+
+    def filter_code(completion: str) -> str:
+        # The program tends to overwrite, we only take the first function
+        completion = completion.lstrip("\n")
+        return completion.split("\n\n")[0]
+
+    def fix_indents(text: str) -> str:
+        return text.replace("\t", "    ")
+
+    return [filter_code(fix_indents(completion)) for completion in batch_completions]
+
+
+def check_correctness(problem: Dict, completion: str, timeout: float,
+                      completion_id: Optional[int] = None) -> Dict:
+    """
+    Evaluates the functional correctness of a completion by running the test
+    suite provided in the problem. 
+
+    :param completion_id: an optional completion ID so we can match
+        the results later even if execution finishes asynchronously.
+    """
+
+    def unsafe_execute():
+        with create_tempdir():
+            # These system calls are needed when cleaning up tempdir.
+            import os
+            import shutil
+            rmtree = shutil.rmtree
+            rmdir = os.rmdir
+            chdir = os.chdir
+
+            # Disable functionalities that can make destructive changes to the test.
+            reliability_guard()
+
+            # Construct the check program and run it.
+            check_program = (
+                problem["prompt"] + completion + "\n" +
+                problem["test"] + "\n" +
+                f"check({problem['entry_point']})"
+            )
+            
+            try:
+                exec_globals = {}
+                with swallow_io():
+                    with time_limit(timeout):
+                        exec(check_program, exec_globals)
+                result.append("passed")
+            except TimeoutException:
+                result.append("timed out")
+            except BaseException as e:
+                result.append(f"failed: {e}")
+
+            # Needed for cleaning up.
+            shutil.rmtree = rmtree
+            os.rmdir = rmdir
+            os.chdir = chdir
+
+    manager = multiprocessing.Manager()
+    result = manager.list()
+
+    p = multiprocessing.Process(target=unsafe_execute)
+    p.start()
+    p.join(timeout=timeout + 1)
+    if p.is_alive():
+        p.kill()
+
+    if not result:
+        result.append("timed out")
+
+    return dict(
+        task_id=problem["task_id"],
+        passed=result[0] == "passed",
+        result=result[0],
+        completion_id=completion_id,
+    )
+
+
+@contextlib.contextmanager
+def time_limit(seconds: float):
+    def signal_handler(signum, frame):
+        raise TimeoutException("Timed out!")
+    signal.setitimer(signal.ITIMER_REAL, seconds)
+    signal.signal(signal.SIGALRM, signal_handler)
+    try:
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+
+
+@contextlib.contextmanager
+def swallow_io():
+    stream = WriteOnlyStringIO()
+    with contextlib.redirect_stdout(stream):
+        with contextlib.redirect_stderr(stream):
+            with redirect_stdin(stream):
+                yield
+
+
+@contextlib.contextmanager
+def create_tempdir():
+    with tempfile.TemporaryDirectory() as dirname:
+        with chdir(dirname):
+            yield dirname
+
+
+class TimeoutException(Exception):
+    pass
+
+
+class WriteOnlyStringIO(io.StringIO):
+    """ StringIO that throws an exception when it's read from """
+
+    def read(self, *args, **kwargs):
+        raise IOError
+
+    def readline(self, *args, **kwargs):
+        raise IOError
+
+    def readlines(self, *args, **kwargs):
+        raise IOError
+
+    def readable(self, *args, **kwargs):
+        """ Returns True if the IO object can be read. """
+        return False
+
+
+class redirect_stdin(contextlib._RedirectStream):  # type: ignore
+    _stream = 'stdin'
+
+
+@contextlib.contextmanager
+def chdir(root):
+    if root == ".":
+        yield
+        return
+    cwd = os.getcwd()
+    os.chdir(root)
+    try:
+        yield
+    except BaseException as exc:
+        raise exc
+    finally:
+        os.chdir(cwd)
+
+def stream_jsonl(filename: str) -> Iterable[Dict]:
+    """
+    Parses each jsonl line and yields it as a dictionary
+    """
+    if filename.endswith(".gz"):
+        with open(filename, "rb") as gzfp:
+            with gzip.open(gzfp, 'rt') as fp:
+                for line in fp:
+                    if any(not x.isspace() for x in line):
+                        yield json.loads(line)
+    else:
+        with open(filename, "r") as fp:
+            for line in fp:
+                if any(not x.isspace() for x in line):
+                    yield json.loads(line)
+
+def estimate_pass_at_k(
+    num_samples: Union[int, List[int], np.ndarray],
+    num_correct: Union[List[int], np.ndarray],
+    k: int,
+) -> np.ndarray:
+    """
+    Estimates pass@k of each problem and returns them in an array.
+    """
+
+    def estimator(n: int, c: int, k: int) -> float:
+        """
+        Calculates 1 - comb(n - c, k) / comb(n, k).
+        """
+        if n - c < k:
+            return 1.0
+        return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))
+
+    if isinstance(num_samples, int):
+        num_samples_it = itertools.repeat(num_samples, len(num_correct))
+    else:
+        assert len(num_samples) == len(num_correct)
+        num_samples_it = iter(num_samples)
+
+    return np.array(
+        [estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)]
+    )
+
+
+def evaluate_functional_correctness(
+    sample_file: str,
+    k: List[int] = [1, 10, 100],
+    n_workers: int = 4,
+    timeout: float = 3.0,
+):
+    problems = {example["task_id"]: example for example in load_dataset("openai_humaneval")["test"]}
+
+    # Check the generated samples against test suites.
+    with ThreadPoolExecutor(max_workers=n_workers) as executor:
+        futures = []
+        completion_id = Counter()
+        n_samples = 0
+        results = defaultdict(list)
+
+        print("Reading samples...")
+        for sample in tqdm(stream_jsonl(sample_file)):
+            task_id = sample["task_id"]
+            completion = sample["completion"]
+            args = (problems[task_id], completion, timeout, completion_id[task_id])
+            future = executor.submit(check_correctness, *args)
+            futures.append(future)
+            completion_id[task_id] += 1
+            n_samples += 1
+
+        if len(completion_id) < len(problems):
+            include_keys = list(problems.keys())[:len(completion_id)]
+            print(f"Only found {len(completion_id)} solutions, reducing problems from {len(problems)}...")
+            problems = {k:v for k,v in problems.items() if k in include_keys}
+
+        assert len(completion_id) == len(problems), "Some problems are not attempted."
+
+        print("Running test suites...")
+        for future in tqdm(as_completed(futures), total=len(futures)):
+            result = future.result()
+            results[result["task_id"]].append((result["completion_id"], result))
+
+    # Calculate pass@k.
+    total, correct = [], []
+    for result in results.values():
+        result.sort()
+        passed = [r[1]["passed"] for r in result]
+        total.append(len(passed))
+        correct.append(sum(passed))
+    total = np.array(total)
+    correct = np.array(correct)
+
+    ks = k
+    pass_at_k = {
+        f"pass@{k}": estimate_pass_at_k(total, correct, k).mean()
+        for k in ks
+        if (total >= k).all()
+    }
+
+    # Finally, save the results in one file:
+    def combine_results():
+        for sample in stream_jsonl(sample_file):
+            task_id = sample["task_id"]
+            result = results[task_id].pop(0)
+            sample["result"] = result[1]["result"]
+            sample["passed"] = result[1]["passed"]
+            yield sample
+
+    return pass_at_k
+
+def reliability_guard(maximum_memory_bytes: Optional[int] = None):
+    """
+    This disables various destructive functions and prevents the generated code
+    from interfering with the test (e.g. fork bomb, killing other processes,
+    removing filesystem files, etc.)
+
+    WARNING
+    This function is NOT a security sandbox. Untrusted code, including, model-
+    generated code, should not be blindly executed outside of one. See the 
+    Codex paper for more information about OpenAI's code sandbox, and proceed
+    with caution.
+    """
+
+    import platform, faulthandler
+
+    if maximum_memory_bytes is not None:
+        import resource
+        resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes))
+        resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes))
+        if not platform.uname().system == 'Darwin':
+            resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))
+
+    faulthandler.disable()
+
+    import builtins
+    builtins.exit = None
+    builtins.quit = None
+
+    import os
+    os.environ['OMP_NUM_THREADS'] = '1'
+
+    os.kill = None
+    os.system = None
+    os.putenv = None
+    os.remove = None
+    os.removedirs = None
+    os.rmdir = None
+    os.fchdir = None
+    os.setuid = None
+    os.fork = None
+    os.forkpty = None
+    os.killpg = None
+    os.rename = None
+    os.renames = None
+    os.truncate = None
+    os.replace = None
+    os.unlink = None
+    os.fchmod = None
+    os.fchown = None
+    os.chmod = None
+    os.chown = None
+    os.chroot = None
+    os.fchdir = None
+    os.lchflags = None
+    os.lchmod = None
+    os.lchown = None
+    os.getcwd = None
+    os.chdir = None
+
+    import shutil
+    shutil.rmtree = None
+    shutil.move = None
+    shutil.chown = None
+
+    import subprocess
+    subprocess.Popen = None  # type: ignore
+
+    import sys
+    sys.modules['ipdb'] = None
+    sys.modules['joblib'] = None
+    sys.modules['resource'] = None
+    sys.modules['psutil'] = None
+    sys.modules['tkinter'] = None
+
+if __name__ == '__main__':
+    os.environ["TOKENIZERS_PARALLELISM"] = "true"
+
+    from awq import AutoAWQForCausalLM
+    model_path = 'TheBloke/zephyr-7B-beta-AWQ'
+    model = AutoAWQForCausalLM.from_quantized(model_path, device_map="auto", max_new_tokens=2048)
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    eval_humaneval(model, tokenizer)

awq/evaluation/kl_divergence.py

+# adapted from https://gist.github.com/Ttl/0d51f739dc59254b4b2183e259c97d82
+
+import torch
+import numpy as np
+from tqdm import tqdm
+from datasets import load_dataset
+from transformers import (
+    PreTrainedTokenizer,
+    PreTrainedModel,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+)
+
+try:
+    from scipy.stats import bayes_mvs
+    from scipy.stats import t as student_t
+    from scipy.stats.mstats import mquantiles_cimj
+    SCIPY_INSTALLED = True
+except:
+    SCIPY_INSTALLED = False
+
+@torch.jit.script
+def rel_entr(x, y):
+    mask = (x > 0) & (y > 0)
+    result = torch.where(mask, x * torch.log(x / y), torch.zeros_like(x))
+    result[(x > 0) & (y <= 0)] = float('inf')
+    return result
+
+def bin_conf(p, n, z):
+    # Binomial distribution confidence bounds
+    # Bayes estimator when p is degenerate
+    if p == 0:
+        p = 1 / (n + 2)
+    if p == 1:
+        p = 1 - 1 / (n + 2)
+    return z * torch.sqrt(p*(1-p)/n)
+
+def eval_kl_divergence(ref_model: PreTrainedModel, eval_model: PreTrainedModel, tokenizer: PreTrainedTokenizer, seqlen: int):
+    if not SCIPY_INSTALLED:
+        raise Exception("SciPy needs to be installed for KL Divergence evaluation: pip install scipy")
+
+    # load dataset
+    data = load_dataset('wikitext', 'wikitext-2-raw-v1', split='test')
+    data = tokenizer("\n\n".join(data['text']), return_tensors='pt')
+    data = data.input_ids.to(ref_model.device)
+
+    n_samples = data.numel() // seqlen
+
+    alpha = 0.01
+    kls = []
+    top1 = 0
+    top5 = 0
+    top10 = 0
+    eval_top5 = 0
+    eval_top10 = 0
+    samples = 0
+    i = 0
+
+    # start eval
+    with tqdm(range(n_samples), desc="KL Div") as progress_bar:
+        for i in progress_bar:
+            start_index = (i * seqlen)
+            end_index = ((i + 1) * seqlen)
+            batch_len = end_index-start_index
+            batch = data[:, start_index:end_index]
+        
+            # get logits
+            with torch.no_grad():
+                y1 = ref_model(batch)[0]
+                y2 = eval_model(batch)[0]
+
+            # kl divergence
+            y1_probs = torch.softmax(y1, dim=-1)
+            y2_probs = torch.softmax(y2, dim=-1)
+            relative_entropy = rel_entr(y1_probs, y2_probs)
+            kl_div = torch.sum(relative_entropy, dim=-1).squeeze(0)
+            kls.append(torch.nan_to_num(kl_div).tolist())
+            
+            # stats
+            eval_argmax = torch.argmax(y2, axis=-1).squeeze(0)
+            ref_argmax = torch.argmax(y1, axis=-1).squeeze(0)
+            eval_part5 = torch.topk(y2, k=5, dim=-1).indices[:, :, -5].squeeze(0)
+            ref_part5 = torch.topk(y1, k=5, dim=-1).indices[:, :, -5].squeeze(0)
+            eval_part10 = torch.topk(y2, k=10, dim=-1).indices[:, :, -10].squeeze(0)
+            ref_part10 = torch.topk(y1, k=10, dim=-1).indices[:, :, -10].squeeze(0)
+            top1 += (eval_argmax == ref_argmax).sum().item()
+            top5 += ((ref_argmax == eval_part5).sum()).item()
+            top10 += ((ref_argmax == eval_part10).sum()).item()
+            eval_top5 += ((eval_argmax == ref_part5).sum()).item()
+            eval_top10 += ((eval_argmax == ref_part10).sum()).item()
+            samples += batch_len
+
+            progress_bar.set_description(
+                f"KL Div: {torch.mean(torch.Tensor(kls)):.4g}, "
+                f"Top 1: {top1 / samples:.4g}, "
+                f"Top 5: {top5 / samples:.4g}, "
+                f"Top 10: {top10 / samples:.4g}"
+            )
+    
+    z = student_t.ppf(1 - alpha/2, samples)
+    m_conf = z*np.sqrt(np.mean([k**2 for k in kls])/len(kls))
+    m, _, __ = bayes_mvs(kls, 1-alpha)
+    q90 = np.quantile(kls, 0.90)
+    q95 = np.quantile(kls, 0.95)
+    q99 = np.quantile(kls, 0.99)
+    q_bounds = mquantiles_cimj(kls, prob=[0.90, 0.95, 0.99])
+
+    print(" -- ")
+    print(" ** Reference model:", ref_model.config.model_type)
+    print(" ** Evaluation model:", eval_model.config.model_type)
+    print(" -- ")
+    print(f" ** KL Divergence: {m[0]:.6g}, [{m[1][0]:.6g} - {m[1][1]:.6g}]")
+    print(f" ** q90: {q90:.4g}, [{q_bounds[0][0]:.4g} - {q_bounds[1][0]:.4g}]")
+    print(f" ** q95: {q95:.4g}, [{q_bounds[0][1]:.4g} - {q_bounds[1][1]:.4g}]")
+    print(f" ** q99: {q99:.4g}, [{q_bounds[0][2]:.4g} - {q_bounds[1][2]:.4g}]")
+    print(f"max: {np.max(kls):.4g}")
+    print(" -- ")
+    print("Reference top token in eval top-n probability:")
+    print(f" ** ref_top1: {top1 / samples:.4g} ± {bin_conf(top1/samples, samples, z):.4g}")
+    print(f" ** ref_top5: {top5 / samples:.4g} ± {bin_conf(top5/samples, samples, z):.4g}")
+    print(f" ** ref_top10: {top10 / samples:4g} ± {bin_conf(top10/samples, samples, z):.4g}")
+    print("Eval top token in reference top-n probability:")
+    print(f" ** eval_top5: {eval_top5 / samples:.4g} ± {bin_conf(eval_top5/samples, samples, z):.4g}")
+    print(f" ** eval_top10: {eval_top10 / samples:4g} ± {bin_conf(eval_top10/samples, samples, z):.4g}")
+
+if __name__ == '__main__':
+    # ref_model_path = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
+    # eval_model_path = "TinyLlama/TinyLlama-1.1B-intermediate-step-1195k-token-2.5T"
+    ref_model_path = eval_model_path = "gpt2"
+
+    tokenizer = AutoTokenizer.from_pretrained(ref_model_path)
+    ref_model = AutoModelForCausalLM.from_pretrained(ref_model_path, device_map="auto")
+    eval_model = AutoModelForCausalLM.from_pretrained(eval_model_path, device_map="auto")
+
+    eval_kl_divergence(ref_model, eval_model, tokenizer, seqlen=1024)
\ No newline at end of file

awq/utils/eval_utils.py

-import torch
-import torch.nn as nn
-from tqdm import tqdm
-from datasets import load_dataset
-
-def evaluate_perplexity(model, tokenizer):
-    def _perplexity(nlls, n_samples, seqlen):
-        return torch.exp(torch.stack(nlls).sum() / (n_samples * seqlen))
-    
-    # load and prepare dataset
-    data = load_dataset('wikitext', 'wikitext-2-raw-v1', split='test')
-    data = tokenizer("\n\n".join(data['text']), return_tensors='pt')
-    data = data.input_ids.to(model.device)
-
-    seqlen = 2048
-    model = model.eval()
-    n_samples = data.numel() // seqlen
-    
-    nlls = []
-
-    with tqdm(range(n_samples), desc="Perplexity -") as progress_bar:
-        for i in progress_bar:
-            start_index = (i * seqlen)
-            end_index = ((i + 1) * seqlen)
-            batch = data[:, start_index:end_index].to(model.device)
-            with torch.no_grad():
-                logits = model(batch).logits
-            shift_logits = logits[:, :-1, :].contiguous().float()
-            shift_labels = data[:, start_index:end_index][:, 1:]
-            loss_fct = nn.CrossEntropyLoss()
-            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-            neg_log_likelihood = loss.float() * seqlen
-            nlls.append(neg_log_likelihood)
-
-            curr_ppl = _perplexity(nlls, i+1, seqlen)
-            progress_bar.set_description(f"Perplexity {curr_ppl:.3f}")
-
-    ppl = _perplexity(nlls, n_samples, seqlen)
-    
-    return ppl.item()
-
-if __name__ == '__main__':
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-    
-    model_path = 'mistralai/Mistral-7B-Instruct-v0.1'
-    model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
-
-    evaluate_perplexity(model, tokenizer)

examples/eval.py

@@ -2,7 +2,13 @@ import argparse
 from lm_eval import evaluator
 from awq import AutoAWQForCausalLM
 from transformers import AutoTokenizer
-from awq.utils.eval_utils import evaluate_perplexity
+from awq.evaluation import (
+    evaluate_perplexity,
+    eval_librispeech,
+    eval_mmlu,
+    eval_humaneval,
+    eval_kl_divergence,
+)
 
 def run_eval(
         model_path, quant_file, device, tasks, task_batch_size, task_n_shot,
@@ -11,18 +17,32 @@ def run_eval(
     """
     Post quantization: Evaluate perplexity on wikitext with EleutherAI Evaluation Harness
     """
+    tasks = tasks.split(',')
+
     # Load model
-    if task_use_pretrained:
-        model = AutoAWQForCausalLM.from_pretrained(model_path, safetensors=pretrained_safetensors)
-    else:
-        model = AutoAWQForCausalLM.from_quantized(model_path, quant_file, fuse_layers=False)
+    if len(tasks) == 1 and tasks[0] != "mmlu" and tasks[0] != "librispeech":
+        if task_use_pretrained:
+            model = AutoAWQForCausalLM.from_pretrained(model_path, safetensors=pretrained_safetensors)
+        else:
+            model = AutoAWQForCausalLM.from_quantized(model_path, quant_file, fuse_layers=False)
 
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
     # Load adapter
-    tasks = tasks.split(',')
     if len(tasks) == 1 and tasks[0] == 'wikitext':
         evaluate_perplexity(model.model, tokenizer)
+    
+    elif len(tasks) == 1 and tasks[0] == 'librispeech':
+        eval_librispeech(model_path)
+    
+    elif len(tasks) == 1 and tasks[0] == 'mmlu':
+        eval_mmlu(model_path, task_n_shot, task_batch_size, device, task_use_pretrained)
+    
+    elif len(tasks) == 1 and tasks[0] == 'humaneval':
+        eval_humaneval(model, tokenizer)
+    
+    elif len(tasks) == 1 and tasks[0] == 'kldiv':
+        eval_kl_divergence(model.model, model.model, tokenizer, seqlen=1024)
 
     else:
         # Evaluate perplexity of quantized model
@@ -43,6 +63,9 @@ if __name__ == '__main__':
 
     - Run perplexity unquantized FP16 model:
     python examples/eval.py --use_pretrained --model_path lmsys/vicuna-7b-v1.5
+
+    - Run MMLU of quantized model:
+    python examples/eval.py --model_path TheBloke/zephyr-7B-beta-AWQ --tasks mmlu --n_shot 1 --batch_size 4
     """
 
     parser = argparse.ArgumentParser()

setup.py

@@ -58,5 +58,14 @@ if platform.system().lower() != "darwin" and HAS_CUDA:
 setup(
     packages=find_packages(),
     install_requires=requirements,
+    extras_require={
+        "eval": [
+            "lm_eval>=0.4.0",
+            "tabulate",
+            "protobuf",
+            "evaluate",
+            "scipy"
+        ],
+    },
     **common_setup_kwargs
 )