Update README.md

README.md

@@ -70,110 +70,6 @@ with torch.no_grad():
     print(out)
 ```
 
-usning [NBCE](https://github.com/bojone/NBCE/tree/main) Inference
-
-```python
-import json, jsonlines
-import torch
-from transformers import AutoTokenizer
-from transformers import AquilaForCausalLM
-from transformers import TopPLogitsWarper, LogitsProcessorList
-
-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.pad_token = tokenizer.unk_token
-
-# load Aquila model
-model = AquilaForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)
-device = torch.device('cuda')
-model.to(device)
-
-# 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, 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)
-    past_key_values = None
-    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,
-                        use_cache=True,
-                        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]
-        logits = logits - logits.logsumexp(dim=-1, keepdims=True)
-        logits = processors(input_ids, logits)
-        entropy = -(logits.exp() * logits.clip(-100, 0)).sum(dim=-1)
-        if i > 0:
-            entropy[k] -= eta
-        k = entropy[1:].argmin() + 1
-        logits_max = logits[k]
-        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 =====
-
-        # 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)
-        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)
-
-
-if __name__ == '__main__':
-    count = 0
-    with open("gaokao_chinese_dataset.jsonl",'r') as f:
-        for item in jsonlines.Reader(f):
-            batch = preprocess(item['prompt'],question=item['question'])
-            generate(10, batch)
-```
 
 ## License