Update README.md

README.md

@@ -73,41 +73,40 @@ with torch.no_grad():
 usning [NBCE](https://github.com/bojone/NBCE/tree/main) Inference
 
 ```python
-import json
+import json, jsonlines
 import torch
 from transformers import AutoTokenizer
-from transformers import AutoModelForCausalLM
+from transformers import AquilaForCausalLM
 from transformers import TopPLogitsWarper, LogitsProcessorList
-import pdb
+
+from cyg_conversation import default_conversation
+def preprocess(text, question="回答："):
+    tmp=""
+    import json
+    contexts = []
+    conv = default_conversation.copy()
+    conv.append_message(conv.roles[0], ""+question)
+    conv.append_message(conv.roles[1], None)
+    contexts.append(conv.get_prompt())
+    for pos in range(0,len(text),1024):
+        conv1 = default_conversation.copy()
+        conv1.append_message(conv1.roles[0], text[pos:min(pos + 1024, len(text))] + question)
+        conv1.append_message(conv1.roles[1], None)
+        contexts.append(conv1.get_prompt())
+    print('Context长度分布：', [len(text) for text in contexts])
+    print('Context总长度：', sum([len(text) for text in contexts]))
+    return contexts
 
 # load tokenizer
+model_path = "checkpoints/hf_weight"
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-tokenizer.padding_side = 'left' 
+tokenizer.padding_side = 'left'
 tokenizer.pad_token = tokenizer.unk_token
 
 # load Aquila model
-model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)
+model = AquilaForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)
 device = torch.device('cuda')
 model.to(device)
-# load example Context
-from cyg_conversation import default_conversation
-
-conv = default_conversation.copy()
-contexts = json.load(open('code_text_2.json'))
-
-question = "请解释这段程序的功能："
-batch = []
-conv.append_message(conv.roles[0], question)
-conv.append_message(conv.roles[1], None)
-batch.append(conv.get_prompt())
-# concat context and question
-for ci,context in enumerate(contexts):        
-    conv1 = default_conversation.copy()
-    conv1.append_message(conv.roles[0], context+question)
-    conv1.append_message(conv.roles[1], None)
-    batch.append(conv1.get_prompt())
-print('Context长度分布：', [len(text) for text in batch])
-print('Context总长度：', sum([len(text) for text in batch]))
 
 # Top-P
 processors = LogitsProcessorList()
@@ -115,17 +114,17 @@ 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 example code
+def generate(max_tokens, batch):
+    """Naive Bayes-based Context Extension 演示代码
     """
     inputs = tokenizer(batch, padding='longest', return_tensors='pt').to(device)
     input_ids = inputs.input_ids
     attention_mask = inputs.attention_mask
-    
-    print('input_ids', input_ids.shape)
+
+    #print('input_ids', input_ids.shape)
     past_key_values = None
     n = input_ids.shape[0]
-    
+
     for i in range(max_tokens):
         # model output
         outputs = model(input_ids=input_ids,
@@ -135,7 +134,7 @@ def generate(max_tokens):
                         past_key_values=past_key_values
                        )
         past_key_values = outputs.past_key_values
-        
+
         # ===== NBCE core code starts =====
         beta, eta = 0.25, 0.1
         logits = outputs.logits[:, -1]
@@ -150,27 +149,30 @@ def generate(max_tokens):
         logits_merged = (1 + beta) * logits_max - beta * logits_uncond
         logits = torch.where(logits_uncond > -100, logits_merged, logits_max)
         # ===== NBCE core code ends =====
-        
+
         # 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)        
+        next_tokens = torch.multinomial(probas, num_samples=1).squeeze(1)
         if next_tokens[0] == tokenizer.eos_token_id:
             break
-            
+
         ret = tokenizer.batch_decode(next_tokens)
         print(ret[0], flush=True, end='')
-        
+
         # prepare for next iteration
         input_ids = next_tokens.unsqueeze(-1).tile(n, 1)
-        attention_mask = torch.cat([attention_mask, torch.ones(n, 1, dtype=torch.long, device=device)], dim=-1)        
+        attention_mask = torch.cat([attention_mask, torch.ones(n, 1, dtype=torch.long, device=device)], dim=-1)
 
 
 if __name__ == '__main__':
-    generate(1000)
-
+    count = 0
+    with open("/data2/gaokao_chinese_dataset.jsonl",'r') as f:
+        for item in jsonlines.Reader(f):
+            batch = preprocess(item['prompt'],question=item['question'])
+            generate(10, batch)
 ```
 
 ## License