MedMemoryBench/main.py at main · AQ-MedAI/MedMemoryBench · 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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
#!/usr/bin/env python3
"""
Personalization Lab - Memory Evaluation System
"""

import argparse
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).parent.resolve()
sys.path.insert(0, str(PROJECT_ROOT))

from src.config import ConfigLoader, PROJECT_ROOT
from src.evaluator import create_evaluator
from src.agent import list_available_methods


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Personalization Lab - Memory Evaluation System",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    parser.add_argument("-m", "--method", type=str, help="Method config name")
    parser.add_argument("-d", "--dataset", type=str, help="Dataset name")
    parser.add_argument("-o", "--output-dir", type=str, default="./outputs", help="Output directory")
    parser.add_argument("--dry-run", action="store_true", help="Dry run mode")
    parser.add_argument("-q", "--quiet", action="store_true", help="Quiet mode")
    parser.add_argument("--list-methods", action="store_true", help="List available methods")
    parser.add_argument("--list-datasets", action="store_true", help="List available datasets")
    parser.add_argument("--list-agents", action="store_true", help="List available agent types")
    parser.add_argument("--resume", action="store_true", help="Resume from checkpoint")

    return parser.parse_args()


def list_methods(config_loader: ConfigLoader) -> None:
    print("\nAvailable methods:")
    print("=" * 60)

    configs = config_loader.list_method_configs()
    if not configs:
        print("  (none)")
    else:
        for name in sorted(configs):
            try:
                cfg = config_loader.load_method_config(name)
                print(f"  {name}")
                print(f"    type: {cfg.method_type}, model: {cfg.model.name}")
            except Exception as e:
                print(f"  {name} (load failed: {e})")

    print("=" * 60)


def list_datasets(config_loader: ConfigLoader) -> None:
    print("\nAvailable datasets:")
    print("=" * 60)

    configs = config_loader.list_dataset_configs()
    if not configs:
        print("  (none)")
    else:
        for name in sorted(configs):
            try:
                cfg = config_loader.load_dataset_config(name)
                print(f"  {name} ({cfg.language})")
            except Exception as e:
                print(f"  {name} (load failed: {e})")

    print("=" * 60)


def list_agents_info() -> None:
    print("\nAvailable agent types:")
    print("=" * 60)
    for method in list_available_methods():
        print(f"  - {method}")
    print("=" * 60)


def main() -> int:
    args = parse_args()
    config_loader = ConfigLoader()

    if args.list_methods:
        list_methods(config_loader)
        return 0

    if args.list_datasets:
        list_datasets(config_loader)
        return 0

    if args.list_agents:
        list_agents_info()
        return 0

    if not args.method:
        print("Error: please specify method (-m/--method)")
        return 1

    if not args.dataset:
        print("Error: please specify dataset (-d/--dataset)")
        return 1

    try:
        evaluator = create_evaluator(
            method_config_name=args.method,
            dataset_name=args.dataset,
            config_loader=config_loader,
            output_dir=Path(args.output_dir),
            dry_run=args.dry_run,
            verbose=not args.quiet,
            resume=args.resume,
        )
    except FileNotFoundError as e:
        print(f"Error: {e}")
        return 1
    except Exception as e:
        print(f"Init failed: {e}")
        return 1

    print(f"\nMethod: {args.method}, Dataset: {args.dataset}, Dry Run: {args.dry_run}\n")

    try:
        report = evaluator.run()
        summary = report.summary
        print(f"\nResults: {summary.get('correct', 0)}/{summary.get('total', 0)} "
              f"({summary.get('overall_accuracy', 0):.2%})")
        print(f"Output: {args.output_dir}")
        return 0
    except KeyboardInterrupt:
        print("\nInterrupted")
        return 130
    except Exception as e:
        print(f"\nFailed: {e}")
        import traceback
        traceback.print_exc()
        return 1


if __name__ == "__main__":
    sys.exit(main())