CSCI2270/query.py at main · spanning-tree/CSCI2270 · 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
import json
import argparse
from typing import List, Dict, Any, Literal

from entity.EntitySRE import EntitySRE
from services.RelationExtractorService import RelationExtractorService
from utils.JsonToCsv import JsonToCSV
from utils.models import NomicEmbedding
from utils.model_mapping import get_llm, get_embedding
from services.SREGeneratorService import SREGenerator

DB_HOST = "localhost"
DB_PORT = 27017
ATLAS_DEPLOYMENT = "myLocalAtlas"

def init_parser():
    parser = argparse.ArgumentParser(description="Extract structured relations using SRE patterns.")
    parser.add_argument("--query", type=str, required=True,
                        help="SRE pattern (e.g. (<ORGAN>)(<SYMPTOM>)(<DIAGNOSIS>))")
    parser.add_argument("--mode", type=str, default="chunk", choices=["chunk", "sentence"],
                        help="Extraction mode - chunk or sentence (default: chunk)")
    parser.add_argument("--db_name", type=str, default="radiology",
                        help="Database name to use")
    parser.add_argument("--output_path", type=str, default="output/query",
                        help="Base path for CSV output file")

    parser.add_argument("--llm_model", type=str, default="gemini-2.5-flash-preview-04-17",
                        help="LLM model to use")
    parser.add_argument("--embedding_model", type=str, default="nomic-embed-text",
                        help="Embedding model to use")

    parser.add_argument("--log_path", type=str, default=None,
                        help="Path to save the log")
    parser.add_argument("--new_log", action="store_true",
                        help="Create a new log file")

    parser.add_argument("--use_async", action="store_true",
                        help="Use async mode")

    return parser.parse_args()

def main():
    args = init_parser()

    pattern = args.query
    mode: Literal["chunk", "sentence"] = args.mode
    db_name = args.db_name

    print(f"Using extraction mode: {mode}")

    llm_model = get_llm(args.llm_model)
    embedding_model = get_embedding(args.embedding_model)

    extractor = RelationExtractorService(
        host=DB_HOST,
        port=DB_PORT,
        db_name=db_name,
        atlas_deployment=ATLAS_DEPLOYMENT,
        llm_model=llm_model,
        embedding_model=embedding_model
    )

    try:
        relations: List[Dict[str, Any]] = extractor.extract_relations(
            EntitySRE(pattern),
            mode=mode,
            use_async=args.use_async
        )["result"]

        relations.sort(key=lambda x: int(x.get('data_row_id', 0)))


        if relations and 'data_row_id' in relations[0]:
            fields = ['data_row_id']
            sample_record = relations[0]
            for field in sample_record.keys():
                if field != 'data_row_id':
                    fields.append(field)

            JsonToCSV.export_to_csv(relations, args.output_path, fields)
        else:
            JsonToCSV.export_to_csv(relations, args.output_path)

        print(f"Results exported to {args.output_path}")

        if args.log_path:
            llm_model.write_log(args.log_path)
        if args.new_log:
            llm_model.write_new_log(f"query/{mode}/")

    finally:
        pass


if __name__ == "__main__":
    main()