Function main

evals/ppl.py:170–225

()

def main():
    parser = argparse.ArgumentParser(description="Evaluate perplexity")
    parser.add_argument('-p', '--path', type=str, default='fla-hub/gla-1.3B-100B')
    parser.add_argument('-d', '--data', type=str, default='fla-hub/slimpajama-test')
    parser.add_argument('-s', '--split', type=str, default='train')
    parser.add_argument('-n', '--column_name', type=str, default='text')
    parser.add_argument('--block_size', type=int, default=28672)
    parser.add_argument('--bucket_size', type=int, default=2048)
    parser.add_argument('--batch_size', type=int, default=1)
    args = parser.parse_args()

    # Set device and random seed
    device = "cuda"
    torch.manual_seed(0)

    # Load model and tokenizer
    print(f"Loading model {args.path}")
    tokenizer = AutoTokenizer.from_pretrained(args.path)
    model = AutoModelForCausalLM.from_pretrained(
        args.path,
        device_map={"": device}
    ).bfloat16().eval()
    print(f"{model}")

    # Load dataset
    print(f"Loading data {args.data}")
    dataset = load_dataset(args.data, split=args.split)
    dataset = dataset.map(
        partial(PerplexityEvaluator.preprocess, tokenizer=tokenizer, column_name=args.column_name),
        batched=True,
        num_proc=32
    )
    print(dataset)
    print("batch_size", args.batch_size, "block_size", args.block_size, "total_tokens_per_batch", args.batch_size * args.block_size)    

    # Create evaluator and run evaluation
    evaluator = PerplexityEvaluator(
        model=model,
        tokenizer=tokenizer,
        device=device,
        block_size=args.block_size,
        bucket_size=args.bucket_size,
        batch_size=args.batch_size
    )
    
    with torch.no_grad():
        results = evaluator.evaluate(dataset)
    
    # Print results
    print("\nEvaluation Results:")
    print(f"Final Perplexity: {results['perplexity']:.2f}")
    print(f"Total Tokens: {results['total_tokens']}")
    print(f"Total Sentences: {results['total_sentences']}")
    print("\nBlock-wise Perplexities:")
    for i, ppl in enumerate(results['block_perplexities']):
        print(f"Block {i}: {ppl:.2f}")

if __name__ == "__main__":