Update llm_utils.py

llm_toolkit/llm_utils.py

@@ -12,6 +12,10 @@ from tqdm import tqdm
 
 def get_template(model_name):
     model_name = model_name.lower()
+    if "mistral" in model_name:
+        return "mistral"
+    if "qwen" in model_name:
+        return "qwen"
     if "llama" in model_name:
         return "llama3"
     if "internlm" in model_name:
@@ -118,44 +122,6 @@ def test_model(model, tokenizer, prompt, device="cuda"):
     )
 
 
-def extract_answer(text, debug=False):
-    if text:
-        # Remove the begin and end tokens
-        text = re.sub(
-            r".*?(assistant|\[/INST\]).+?\b",
-            "",
-            text,
-            flags=re.DOTALL | re.MULTILINE,
-        )
-        if debug:
-            print("--------\nstep 1:", text)
-
-        text = re.sub(r"<.+?>.*", "", text, flags=re.DOTALL | re.MULTILINE)
-        if debug:
-            print("--------\nstep 2:", text)
-
-        text = re.sub(
-            r".*?end_header_id\|>\n\n", "", text, flags=re.DOTALL | re.MULTILINE
-        )
-        if debug:
-            print("--------\nstep 3:", text)
-
-        text = text.split("。")[0].strip()
-        if debug:
-            print("--------\nstep 4:", text)
-
-        text = re.sub(
-            r"^Response:.+?\b",
-            "",
-            text,
-            flags=re.DOTALL | re.MULTILINE,
-        )
-        if debug:
-            print("--------\nstep 5:", text)
-
-    return text
-
-
 def eval_model(
     model,
     tokenizer,