diff --git a/python/llm/dev/benchmark/perplexity/README.md b/python/llm/dev/benchmark/perplexity/README.md
index bcf42ff4..870c8dd8 100644
--- a/python/llm/dev/benchmark/perplexity/README.md
+++ b/python/llm/dev/benchmark/perplexity/README.md
@@ -1,25 +1,32 @@
 # Perplexity
-Perplexity (PPL) is one of the most common metrics for evaluating language models. This benchmark implementation was from [transformers/perplexity](https://huggingface.co/docs/transformers/perplexity#perplexity-of-fixed-length-models) and [benchmark_patch_llm.py](https://github.com/insuhan/hyper-attn/blob/main/benchmark_patch_llm.py) 
+Perplexity (PPL) is one of the most common metrics for evaluating language models. This benchmark implementation is adapted from [transformers/perplexity](https://huggingface.co/docs/transformers/perplexity#perplexity-of-fixed-length-models) and [benchmark_patch_llm.py](https://github.com/insuhan/hyper-attn/blob/main/benchmark_patch_llm.py) 
+
+## Run on Wikitext
+
+Download the dataset from [here](https://paperswithcode.com/dataset/wikitext-2), unzip it and we will use the test dataset `wiki.test.raw` for evaluation.
 
-## HOW TO RUN
 ```bash
-python run.py --model_path <path/to/model> --precisions sym_int4 fp4 mixed_fp4 sym_int8 fp8_e5m2 fp8_e4m3 mixed_fp8 --device xpu --datasets dataset_names --dataset_path <path/to/dataset> --language en
+python run_wikitext.py --model_path meta-llama/Meta-Llama-3-8B/ --data_path wikitext-2-raw-v1/wikitext-2-raw/wiki.test.raw --precision sym_int4 --use-cache --device xpu
+
+# Run with stride
+python run_wikitext.py --model_path meta-llama/Meta-Llama-3-8B/ --data_path wikitext-2-raw-v1/wikitext-2-raw/wiki.test.raw --precision fp16 --device xpu --stride 512
+```
+
+## Run on [THUDM/LongBench](https://github.com/THUDM/LongBench) dataset
+
+```bash
+python run.py --model_path <path/to/model> --precisions sym_int4 fp8 --device xpu --datasets dataset_names --dataset_path <path/to/dataset> --language en
 ```
 A more specific example to run perplexity on Llama2-7B using the default English datasets:
 ```bash
 python run.py --model_path meta-llama/Llama-2-7b-chat-hf --precisions float16 sym_int4 --device xpu --language en
 ```
 
-> Note: We currently only support the `THUDM/LongBench` [dataset](https://github.com/THUDM/LongBench)
-
+Notes:
 - If you want to test model perplexity on a few selected datasets from the `LongBench` dataset, please use the format below.
   ```bash
   --datasets narrativeqa qasper ...
   ```
 - The `language` argument will only take effect if `datasets` is `None`. The choices for this argument are `en, zh, all`, which stands for all the English datasets, all the Chinese datasets and all the datasets respectively during testing.
 - If you want to test perplexity on pre-downloaded datasets, please specify the `<path/to/dataset>` in the `dataset_path` argument in your command.
-
-## Summarize the results
-```python
-python make_table.py <input_dir>
-```
+- You can run `python make_table.py <input_dir>` to summarize the results.
diff --git a/python/llm/dev/benchmark/perplexity/ppl.py b/python/llm/dev/benchmark/perplexity/ppl.py
index d340297e..b0e4ac42 100644
--- a/python/llm/dev/benchmark/perplexity/ppl.py
+++ b/python/llm/dev/benchmark/perplexity/ppl.py
@@ -13,6 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+# This file is adapted from
+# https://github.com/insuhan/hyper-attn/blob/main/benchmark_patch_llm.py
+#
 
 import numpy as np
 import torch
diff --git a/python/llm/dev/benchmark/perplexity/run.py b/python/llm/dev/benchmark/perplexity/run.py
index 881d82e0..92b4999a 100644
--- a/python/llm/dev/benchmark/perplexity/run.py
+++ b/python/llm/dev/benchmark/perplexity/run.py
@@ -13,6 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+# This file is adapted from
+# https://github.com/insuhan/hyper-attn/blob/main/benchmark_patch_llm.py
+#
 
 import argparse
 from tqdm import tqdm
diff --git a/python/llm/dev/benchmark/perplexity/run_wikitext.py b/python/llm/dev/benchmark/perplexity/run_wikitext.py
index f308ed93..a5b7dea3 100644
--- a/python/llm/dev/benchmark/perplexity/run_wikitext.py
+++ b/python/llm/dev/benchmark/perplexity/run_wikitext.py
@@ -1,21 +1,35 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file is adapted from
+# https://huggingface.co/docs/transformers/en/perplexity
+#
+
 import argparse
 import torch
-from datasets import load_dataset
 from tqdm import tqdm
 
 
-def parse_kwargs(kwstr):
-    kvpair = [item.split('=') for item in kwstr.split(',') if item != ""]
-    return {k:v for k, v in kvpair}
-
-
 parser = argparse.ArgumentParser()
 parser.add_argument("--model_path", required=True, type=str)
-parser.add_argument("--dataset", type=str, default='path=wikitext,name=wikitext-2-raw-v1')
+parser.add_argument("--data_path", type=str, default='wikitext-2-raw-v1/wikitext-2-raw/wiki.test.raw')
+parser.add_argument("--chunk_size", type=int, default=512)
+parser.add_argument("--stride", type=int, default=0)
 parser.add_argument("--device", type=str, default="xpu")
 parser.add_argument("--precision", type=str, default="sym_int4")
 parser.add_argument("--use-cache", action="store_true")
-parser.add_argument("--limit", type=int, default=None, help="Limit the number of examples per task. For debug only")
 args = parser.parse_args()
 
 if args.precision == "fp16":  # ipex fp16
@@ -31,24 +45,30 @@ else:  # ipex-llm
     model = model.half()
 model = model.to(args.device)
 
+with open(args.data_path, "rb") as f:
+    data = f.read()
+
 from transformers import AutoTokenizer
 tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
-
-test = load_dataset(**parse_kwargs(args.dataset), split="test")["text"]
-if args.limit:
-    test = test[:args.limit]
-encodings = tokenizer("\n\n".join(test), return_tensors="pt")
+encodings = tokenizer(data.decode("utf-8").strip("\n"), return_tensors="pt")
 
 max_length = model.config.max_position_embeddings
-stride = 512
+stride = args.chunk_size if args.stride <= 0 else args.stride
 seq_len = encodings.input_ids.size(1)
+num_chunks = seq_len // stride
 
 nlls = []
 prev_end_loc = 0
-for begin_loc in tqdm(range(0, seq_len, stride)):
-    end_loc = min(begin_loc + max_length, seq_len)
-    trg_len = end_loc - prev_end_loc  # may be different from stride on last loop
+for i in tqdm(range(num_chunks)):
+    begin_loc = i * stride
+    if args.stride > 0:
+        end_loc = min(begin_loc + max_length, seq_len)
+        trg_len = end_loc - prev_end_loc  # may be different from stride on last loop
+    else:
+        end_loc = begin_loc + stride
+        trg_len = -stride//2
     input_ids = encodings.input_ids[:, begin_loc:end_loc].to(args.device)
+    if args.stride == 0: input_ids[:, 0] = tokenizer.bos_token_id
     target_ids = input_ids.clone()
     target_ids[:, :-trg_len] = -100
 
@@ -69,4 +89,4 @@ for begin_loc in tqdm(range(0, seq_len, stride)):
         break
 
 ppl = torch.exp(torch.stack(nlls).mean())
-print("ppl result: {}".format(ppl.item()))
+print("Final ppl estimate: {}".format(ppl.item()))