ConflictHarm/config.py at master · DataArcTech/ConflictHarm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class Model_Config():
    def __init__(self):
        self.chat_template_system = []
        self.max_new_tokens = 0
        self.model_infer_kwargs = {}
        self.padding = True
        self.truncation = True
        self.run_api = False
        self.add_additional_assistant_prompt = False
        self.additional_prompt = ""
        self.call_by_api_ = False
        self.use_vllm = False
        self.n_gpus = 0

    def model_setting(self, model_name):
        if 'qwq' in model_name.lower():
            self.max_new_tokens = 32768
            # self.model_infer_kwargs["temperature"]=0.6
        elif 'nemotron' in model_name.lower():
            self.max_new_tokens = 32768
            self.model_infer_kwargs["temperature"]=0.6
            self.model_infer_kwargs["top_p"]=0.95
            thinking = "on"
            self.chat_template_system.append({"role": "system", "content": f"detailed thinking {thinking}"})
            self.additional_prompt = "<think>\n"
            self.add_additional_assistant_prompt = True
        elif 'star1-r1' in model_name.lower():
            self.max_new_tokens = 16384
            self.model_infer_kwargs["temperature"]=0.6
            self.model_infer_kwargs["top_p"]=0.95
            self.additional_prompt = "<think>\n"
        elif 'tars' in model_name.lower():
            self.max_new_tokens = 131072
            self.model_infer_kwargs["temperature"]=0.7
            self.model_infer_kwargs["top_p"]=0.8
            self.additional_prompt = "<think>\n"
        elif 'realsafe-r1' in model_name.lower():
            self.max_new_tokens = 4096
            self.model_infer_kwargs["temperature"]=0.6
            self.model_infer_kwargs["top_p"]=0.95
            self.additional_prompt = "<think>\n"
        elif 'qwen2.5-32b' in model_name.lower():
            self.max_new_tokens = 2048
        elif 'deepseek-r1' in model_name.lower():
            self.model_infer_kwargs["temperature"] = 0
            self.call_by_api_ = True
        else:
            raise KeyError("Model name error.")