Update README.md

README.md

@@ -99,14 +99,32 @@ def special_encode(input, tokenizer):
             res_id.append(sep_id)
 
     return res_id
+
+def special_encode(input, tokenizer):
+    raw_str = "[USER]%s[SEP][BOT]" % input.strip().replace("\r", "")
+    eos_id = tokenizer.eos_token_id
+    bos_id = tokenizer.bos_token_id
+    sep_id = tokenizer.encode("[SEP]")[-1]
+    res_id = [eos_id, bos_id]
+    arr = raw_str.split("[SEP]")
+    for elem_idx in range(len(arr)):
+        elem = arr[elem_idx]
+        elem_id = tokenizer.encode(elem)[1:]
+        res_id += elem_id
+        if elem_idx < len(arr) - 1:
+            res_id.append(sep_id)
+
+    return res_id
+
 if __name__ == '__main__':
     text = "小王要将150千克含药量20%的农药稀释成含药量5%的药水．需要加水多少千克？"
     text_token_ids = torch.tensor(special_encode(
         text, tokenizer)).to(model.device).reshape(1, -1)
     response = model.generate(text_token_ids, do_sample=False, max_length=512)
-    response_text = tokenizer.decode(response.cpu()[0], skip_special_tokens=True).split(
-        "[BOT]")[-1].split("[SEP]")[0].strip()
-    print(response_text)   
+    response_text = tokenizer.decode(response.cpu()[0], skip_special_tokens=True)
+    
+    response_text = extract_res(response_text)
+    print(response_text)    
     """输出结果：
     首先，我们需要计算出150千克含药量20%的农药中含有多少千克的药。\n\n150千克 * 20% = 30千克\n\n然后，我们需要计算出要得到含药量5%的药水，需要多少千克的药水。\n\n30千克 / 5% = 600千克\n\n最后，我们需要计算出需要加多少千克的水。\n\n600千克 - 150千克 = 450千克\n\n所以答案是，小王需要加450千克的水。
     """ 
@@ -141,13 +159,25 @@ def special_encode(input, tokenizer):
             res_id.append(sep_id)
 
     return res_id
+
+def extract_res(response):
+    if "[BOT]" in response:
+        response = response.split("[BOT]")[1]
+    if "<s>" in response:
+        response = response.split("<s>")[-1]
+    if "</s>" in response:
+        response = response.split("</s>")[0]
+    if "[SEP]" in response:
+        response = response.split("[SEP]")[0]
+    return response
+
 if __name__ == '__main__':
     text="Janet’s ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?"
     text_token_ids = torch.tensor(special_encode(
         text, tokenizer)).to(model.device).reshape(1, -1)
     response = model.generate(text_token_ids, do_sample=False, max_length=512)
-    response_text = tokenizer.decode(response.cpu()[0], skip_special_tokens=True).split(
-        "[BOT]")[-1].split("[SEP]")[0].strip()
+    response_text = tokenizer.decode(response.cpu()[0], skip_special_tokens=True)
+    response_text = extract_res(response_text)
     print(response_text)    
     """Skywork-13B-Math Response:
     First, we need to find out how many eggs Janet has left after eating for breakfast and baking for her friends. \n\nShe has 16 eggs per day, eats 3 for breakfast and uses 4 for baking. So, 16 - 3 - 4 = 9 eggs are left for selling at the farmers' market.\n\nSince she sells each egg for $2, she makes 9 * 2 = $<<9*2=18>>18 every day at the farmers' market.\n\nSo, the answer is $18.