Update README.md

README.md

@@ -80,16 +80,16 @@ from transformers import AutoModelForCausalLM
 from transformers import TopPLogitsWarper, LogitsProcessorList
 import pdb
 
-# 加载tokenizer
+# load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 tokenizer.padding_side = 'left' 
 tokenizer.pad_token = tokenizer.unk_token
 
-# 加载Aquila模型
+# load Aquila model
 model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)
 device = torch.device('cuda')
 model.to(device)
-# 加载示例Context
+# load example Context
 from cyg_conversation import default_conversation
 
 conv = default_conversation.copy()
@@ -100,7 +100,7 @@ batch = []
 conv.append_message(conv.roles[0], question)
 conv.append_message(conv.roles[1], None)
 batch.append(conv.get_prompt())
-# 拼接context和question
+# concat context and question
 for ci,context in enumerate(contexts):        
     conv1 = default_conversation.copy()
     conv1.append_message(conv.roles[0], context+question)
@@ -109,14 +109,14 @@ for ci,context in enumerate(contexts):
 print('Context长度分布：', [len(text) for text in batch])
 print('Context总长度：', sum([len(text) for text in batch]))
 
-# Top-P截断
+# Top-P
 processors = LogitsProcessorList()
 processors.append(TopPLogitsWarper(0.95))
 
 # Copied from https://github.com/bojone/NBCE/blob/main/test.py#L51-L106
 @torch.inference_mode()
 def generate(max_tokens):
-    """Naive Bayes-based Context Extension 演示代码
+    """Naive Bayes-based Context Extension example code
     """
     inputs = tokenizer(batch, padding='longest', return_tensors='pt').to(device)
     input_ids = inputs.input_ids
@@ -127,7 +127,7 @@ def generate(max_tokens):
     n = input_ids.shape[0]
     
     for i in range(max_tokens):
-        # 模型输出
+        # model output
         outputs = model(input_ids=input_ids,
                         attention_mask=attention_mask,
                         return_dict=True,
@@ -136,7 +136,7 @@ def generate(max_tokens):
                        )
         past_key_values = outputs.past_key_values
         
-        # ===== 核心代码开始 =====
+        # ===== NBCE core code starts =====
         beta, eta = 0.25, 0.1
         logits = outputs.logits[:, -1]
         logits = logits - logits.logsumexp(dim=-1, keepdims=True)
@@ -149,11 +149,11 @@ def generate(max_tokens):
         logits_uncond = logits[0]
         logits_merged = (1 + beta) * logits_max - beta * logits_uncond
         logits = torch.where(logits_uncond > -100, logits_merged, logits_max)
-        # ===== 核心代码结束 =====
+        # ===== NBCE core code ends =====
         
-        # 构建分布，采样
-        # tau = 1是标准的随机采样，tau->0则是贪心搜索
-        # 简单起见，这里没有实现topk、topp截断
+        # Building a distribution and sampling
+        # tau = 1 is standard random sampling，tau->0 is greedy search
+        # For simplicity, top-k and top-p truncation are not implemented here.
         tau = 0.01
         probas = torch.nn.functional.softmax(logits[None] / tau , dim=-1)
         next_tokens = torch.multinomial(probas, num_samples=1).squeeze(1)