src/kg/main.py

# src.kg.main.py
import os
import json
from pathlib import Path

from .preprocess import preprocess
from .save_triples import save_triples_for_scripts
from .generate_kg import init_kg, refine_kg

def script2kg(scene_list, idx, name, api_key, model_id):
    # 1) preprocess script
    preprocessed_script = preprocess(scene_list, idx)
    
    # 2) extract triples
    triple_list = save_triples_for_scripts(preprocessed_script, idx, api_key, model_id)

    # 3) build kg
    kg = init_kg(preprocessed_script, idx, triple_list)

    # 4) refine kg
    kg = refine_kg(kg, idx, topk=10)
    
    return kg