Update README.md

README.md

@@ -120,21 +120,17 @@ def special_encode(input, tokenizer):
 
     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)
+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
 
-    return res_id
 
 if __name__ == '__main__':
     text = "小王要将150千克含药量20%的农药稀释成含药量5%的药水．需要加水多少千克？"