Plug-in-Expert/main.py at flexlora-dev · Heeqw/Plug-in-Expert · 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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
"""
Main entry point for FlexLoRA.
"""

from classification.scripts.train_lora import train_classification_lora
from classification.scripts.infer_lora import infer_classification_lora
from classification.scripts.train_sft import train_classification_sft
from classification.scripts.infer_sft import infer_classification_sft
from generation.scripts.train_lora_gen_multiclasses import train_lora_gen_multiclasses
from generation.scripts.infer_lora_gen_hardrouter import infer_lora_gen_hardrouter
from generation.scripts.infer_lora_gen_softrouter import infer_lora_gen_softrouter
from generation.scripts.infer_lora_gen_majorvoting import infer_lora_gen_major_voting

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="FlexLoRA: A flexible QLoRA framework")
    parser.add_argument("--mode", choices=["train", "infer"], default="train",
                        help="Mode to run: train, infer")
    parser.add_argument("--task", choices=["Classification","Generation"], default="Generation",
                        help="Task to run: classification, generation")

    # 解析基本参数
    args, remaining = parser.parse_known_args()

    if args.task == "Classification":
        classification_parser = argparse.ArgumentParser()
        classification_parser.add_argument("--training_method", choices=["sft", "lora", "qlora"], default="lora", help="Training method for generation task: sft, lora, or qlora")
        classification_args = classification_parser.parse_args(remaining)
        args.training_method = classification_args.training_method

        if args.mode == "train":
            if args.training_method == "qlora":
                train_classification_lora(use_qlora=True)
            elif args.training_method == "lora":
                train_classification_lora(use_qlora=False)
            elif args.training_method == "sft":
                train_classification_sft()

        elif args.mode == "infer":
            if args.training_method == "qlora":
                infer_classification_lora(use_qlora=True, training_method=args.training_method)
            elif args.training_method == "lora":
                infer_classification_lora(use_qlora=False, training_method=args.training_method)
            elif args.training_method == "sft":
                infer_classification_sft()


    if args.task == "Generation":
        if args.mode == "train":
            train_lora_gen_multiclasses()
        elif args.mode == "infer":
            generation_parser = argparse.ArgumentParser()
            generation_parser.add_argument("--inference_method",
                choices=["hard_route", "soft_route"],
                default="hard_route",
                help="Inference method for generation task: hard_route, soft_route")
            generation_parser.add_argument("--soft_route_method",
                choices=["major_voting", "prior&confidence_fusion"],
                default="major_voting",
                help="Soft route method for generation task: major_voting, soft_voting")
            generation_args = generation_parser.parse_args(remaining)
            args.inference_method = generation_args.inference_method
            args.soft_route_method = generation_args.soft_route_method

            if args.inference_method == "hard_route":
                infer_lora_gen_hardrouter()
            elif args.inference_method == "soft_route":
                if args.soft_route_method == "soft_voting":
                    infer_lora_gen_softrouter()
                elif args.soft_route_method == "major_voting":
                    infer_lora_gen_major_voting()