first commit

.gitignore

@@ -0,0 +1,20 @@
+# Python
+__pycache__/
+*.py[cod]
+*.pyo
+*.egg
+*.egg-info/
+dist/
+build/
+*.log
+
+.env
+*.env.*
+
+*.pem
+*.key
+
+.cache/
+*.pytest_cache/
+
+.gradio/

app.py

@@ -0,0 +1,543 @@
+import gradio as gr
+import json
+from tqdm import tqdm
+import numpy as np
+import random
+import torch
+import ast
+from difflib import HtmlDiff
+
+from src.kg.main import script2kg
+from src.summary.summarizer import Summarizer
+from src.summary.utils import preprocess_script, chunk_script_gpt
+from src.summary.prompt import build_summarizer_prompt
+from src.fact.narrativefactscore import NarrativeFactScore
+
+def _set_seed(seed):
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+def parse_scenes(scene_text):
+    try:
+        return json.loads(scene_text)
+    except json.JSONDecodeError:
+        return ast.literal_eval(scene_text)
+
+def set_name_list(dataset, data_type):
+    if dataset == "MovieSum":
+        if data_type == "train":
+            return ['8MM_1999', 'The Iron Lady_2011', 'Adventureland_2009', 'Napoleon_2023', 
+                   'Kubo and the Two Strings_2016', 'The Woman King_2022', 'What They Had_2018', 
+                   'Synecdoche, New York_2008', 'Black Christmas_2006', 'Superbad_2007']
+        elif data_type == "validation":
+            return ['The Boondock Saints_1999', 'The House with a Clock in Its Walls_2018',
+                   'The Unbelievable Truth_1989', 'Insidious_2010', 'If Beale Street Could Talk_2018',
+                   'The Battle of Shaker Heights_2003', '20th Century Women_2016', 
+                   'Captain Phillips_2013', 'Conspiracy Theory_1997', 'Domino_2005']
+        elif data_type == "test":
+            # Return test dataset names (shortened for brevity)
+            return ['A Nightmare on Elm Street 3: Dream Warriors_1987', 'Van Helsing_2004', 
+                   'Oppenheimer_2023', 'Armored_2009', 'The Martian_2015']
+    elif dataset == "MENSA":
+        if data_type == "train":
+            return ['The_Ides_of_March_(film)', 'An_American_Werewolf_in_Paris',
+                   'Batman_&_Robin_(film)', 'Airplane_II:_The_Sequel', 'Krull_(film)']
+        elif data_type == "validation":
+            return ['Pleasantville_(film)', 'V_for_Vendetta_(film)',
+                   'Mary_Shelleys_Frankenstein_(film)', 'Rapture_(1965_film)', 'Get_Out']
+        elif data_type == "test":
+            return ['Knives_Out', 'Black_Panther', 'Pet_Sematary_(film)',
+                   'Panic_Room', 'The_Village_(2004_film)']
+    return []
+
+def update_name_list_interface(dataset, data_type):
+    if dataset in ["MovieSum", "MENSA"]:
+        return (
+            gr.update(choices=set_name_list(dataset, data_type), value=None, visible=True),
+            gr.update(visible=False),
+            gr.update(value="")
+        )
+    else:
+        return (
+            gr.update(visible=False),
+            gr.update(visible=True),
+            gr.update(value="Click next 'Knowledge Graph' to continue")
+        )
+
+def read_data(dataset, data_type):
+    file_path = f"dataset/{dataset}/{data_type}.jsonl"
+    try:
+        with open(file_path, 'r', encoding='utf8') as f:
+            data = [json.loads(line) for line in f]
+        return data
+    except FileNotFoundError:
+        return []
+
+def find_work_index(data, work_name):
+    for idx, entry in enumerate(data):
+        if entry.get("name") == work_name:
+            return idx, entry
+    return None, "Work not found in the selected dataset."
+
+def get_narrative_content(dataset, data_type, work):
+    data = read_data(dataset, data_type)
+    for entry in data:
+        if entry.get("name") == work:
+            return entry['scenes']
+    return "Work not found in the selected dataset."
+
+def get_narrative_content_with_index(dataset, data_type, work):
+    data = read_data(dataset, data_type)
+    for idx, entry in enumerate(data):
+        if entry.get("name") == work:
+            # For MovieSum and MENSA datasets, only return scenes
+            if dataset in ["MovieSum", "MENSA"]:
+                return "\n".join(entry['scenes']), idx, data
+            # For other datasets or custom input, return full content
+            return entry, idx, data
+    return "Work not found in the selected dataset.", None, None
+
+def show_diff(original, revised):
+    d = HtmlDiff()
+    original_lines = original.splitlines(keepends=True)
+    revised_lines = revised.splitlines(keepends=True)
+    diff_table = d.make_table(original_lines, revised_lines, fromdesc='Original Summary', todesc='Refined Summary', context=True, numlines=2)
+    return diff_table
+
+def extract_initial_summary(summary_result):
+    return summary_result['summary_agg']['summaries']
+
+def extract_factuality_score_and_details(fact_score_result):
+    factuality_score = fact_score_result['fact_score']
+    feedback_list = []
+    for i, feedback_data in enumerate(fact_score_result['summary_feedback_pairs']):
+        feedbacks = [fb for fb in feedback_data['feedbacks'] if fb.strip()]
+        if feedbacks:
+            feedback_list.append(f"In chunk {i + 1}: {'; '.join(feedbacks)}")
+    incorrect_details = "\n".join(feedback_list)
+    return factuality_score, incorrect_details
+
+def build_kg(script, idx, api_key, model_id):
+    kg = script2kg(script['scenes'], idx, script['name'], api_key, model_id)
+    return kg
+
+def build_kg_custom(scenes, idx, api_key, model_id):
+    kg = script2kg(scenes, idx, "custom", api_key, model_id)
+    return kg
+
+def build_kg_with_data(data, work_index, custom_scenes, api_key, model_id):
+    if data and work_index is not None:  # Dataset mode
+        script = data[int(work_index)]
+        try:
+            kg = script2kg(script['scenes'], int(work_index), script['name'], api_key, model_id)
+            return kg, "Knowledge Graph built successfully!"
+        except Exception as e:
+            return None, f"Error building knowledge graph: {str(e)}"
+    elif custom_scenes:  # Custom script mode
+        try:
+            scenes = parse_scenes(custom_scenes)
+            if not isinstance(scenes, list):
+                return None, "Invalid format. Please provide scenes as a list."
+            kg = build_kg_custom(scenes, 0, api_key, model_id)
+            return kg, "Knowledge Graph built successfully!"
+        except (json.JSONDecodeError, SyntaxError, ValueError) as e:
+            return None, f"Invalid format. Error: {str(e)}"
+        except Exception as e:
+            return None, f"Error building knowledge graph: {str(e)}"
+    return None, "Please select a work or input custom scenes."
+
+def generate_summary(script, idx, api_key, model_id):
+    _set_seed(42)
+    scripty_summarizer = Summarizer(
+        inference_mode="org",
+        model_id=model_id,
+        api_key=api_key,
+        dtype="float16",
+        seed=42
+    )
+    scenes = [f"s#{i}\n{s}" for i, s in enumerate(script['scenes'])]
+    script = "\n\n".join(scenes)
+    script_chunks = chunk_script_gpt(script=script, model_id=model_id, chunk_size=2048)
+
+    script_summaries = []
+    for chunk in tqdm(script_chunks):
+        chunk = preprocess_script(chunk)
+        prompt = build_summarizer_prompt(
+            prompt_template="./templates/external_summary.txt",
+            input_text_list=[chunk]
+        )
+        script_summ = scripty_summarizer.inference_with_gpt(prompt=prompt)
+        script_summaries.append(script_summ.strip())
+
+    elem_dict_list = []
+    agg_dict = {
+        'script': ' '.join(script_chunks),
+        'summaries': ' '.join(script_summaries)
+    }
+
+    for i, (chunk, summary) in enumerate(zip(script_chunks, script_summaries)):
+        elem_dict = {
+            "chunk_index": i,
+            "chunk": chunk.strip(),
+            "summary": summary.strip()
+        }
+        elem_dict_list.append(elem_dict)
+
+    processed_dataset = {
+        "script": script,
+        "scenes": scenes,
+        "script_chunks": script_chunks,
+        "script_summaries": script_summaries,
+    }
+
+    return {"summary_sep": elem_dict_list, "summary_agg": agg_dict, "processed_dataset": processed_dataset}
+
+def generate_summary_with_data(data, work_index, custom_scenes, api_key, model_id):
+    if data and work_index is not None:  # Dataset mode
+        script = data[int(work_index)]
+        try:
+            summary = generate_summary(script, int(work_index), api_key, model_id)
+            return summary, extract_initial_summary(summary)
+        except Exception as e:
+            return None, f"Error generating summary: {str(e)}"
+    elif custom_scenes:  # Custom script mode
+        try:
+            scenes = parse_scenes(custom_scenes)
+            if not isinstance(scenes, list):
+                return None, "Invalid format. Please provide scenes as a list."
+            script = {"name": "custom", "scenes": scenes}
+            summary = generate_summary(script, 0, api_key, model_id)
+            return summary, extract_initial_summary(summary)
+        except (json.JSONDecodeError, SyntaxError, ValueError) as e:
+            return None, f"Invalid format. Error: {str(e)}"
+        except Exception as e:
+            return None, f"Error generating summary: {str(e)}"
+    return None, "Please select a work or input custom scenes."
+
+def calculate_narrative_fact_score(summary, kg_raw, api_key, model_id):
+    _set_seed(42)
+    factscorer = NarrativeFactScore(split_type='gpt', model='gptscore', api_key=api_key, model_id=model_id)
+
+    summary = summary['processed_dataset']
+    chunks, summaries = summary['script_chunks'], summary['script_summaries']
+    total_output = {'fact_score': 0, 'summary_feedback_pairs': []}
+    partial_output = {'fact_score': 0, 'summary_feedback_pairs': []}
+    total_score = 0
+    kg = []
+    for elem in kg_raw:
+        if elem['subject'] == elem['object']:
+            kg.append(f"{elem['subject']} {elem['predicate']}")
+        else:
+            kg.append(f"{elem['subject']} {elem['predicate']} {elem['object']}")
+
+    scores, scores_per_sent, relevant_scenes, summary_chunks, feedbacks = factscorer.score_src_hyp_long(chunks, summaries, kg)
+    for i, score in enumerate(scores):
+        output_elem = {
+            'src': chunks[i],
+            'summary': summaries[i],
+            'score': score,
+            'scores_per_sent': scores_per_sent[i],
+            'relevant_scenes': relevant_scenes[i],
+            'summary_chunks': summary_chunks[i],
+            'feedbacks': feedbacks[i],
+        }
+        output_elem_part = {
+            'scores_per_sent': scores_per_sent[i],
+            'summary_chunks': summary_chunks[i],
+            'feedbacks': feedbacks[i],
+        }
+        total_output['summary_feedback_pairs'].append(output_elem)
+        partial_output['summary_feedback_pairs'].append(output_elem_part)
+        total_score += score
+
+    total_output['fact_score'] = float(total_score / len(scores))
+    partial_output['fact_score'] = float(total_score / len(scores))
+    return total_output, partial_output
+
+def refine_summary(summary, fact_score, api_key, model_id):
+    _set_seed(42)
+    threshold = 0.9
+    summarizer = Summarizer(
+        inference_mode="org",
+        model_id=model_id,
+        api_key=api_key,
+        dtype="float16",
+        seed=42
+    )
+
+    processed_dataset = {
+        "script": summary["script"],
+        "scenes": summary["scenes"],
+        "script_chunks": [],
+        "script_summaries": []
+    }
+    elem_dict_list = []
+    agg_dict = {}
+
+    for factscore_chunk in tqdm(fact_score['summary_feedback_pairs']):
+        src_chunk = factscore_chunk['src']
+        original_summary = factscore_chunk['summary']
+
+        if factscore_chunk['score'] >= threshold:
+            processed_dataset["script_chunks"].append(src_chunk)
+            processed_dataset["script_summaries"].append(original_summary.strip())
+            continue
+
+        hallu_idxs = np.where(np.array(factscore_chunk['scores_per_sent']) == 0)[0]
+        hallu_summary_parts = np.array(factscore_chunk['summary_chunks'])[hallu_idxs]
+        feedbacks = np.array(factscore_chunk['feedbacks'])[hallu_idxs]
+
+        prompt = build_summarizer_prompt(
+            prompt_template="./templates/self_correction.txt",
+            input_text_list=[src_chunk, original_summary]
+        )
+
+        for j, (hallu_summ, feedback) in enumerate(zip(hallu_summary_parts, feedbacks)):
+            prompt += f"\n- Statement to Revise {j + 1}: {hallu_summ} (Reason for Revision: {feedback})"
+        prompt += "\n- Revised Summary: "
+
+        revised_summary = summarizer.inference_with_gpt(prompt=prompt)
+
+        if len(revised_summary.strip()) == 0:
+            revised_summary = original_summary
+
+        processed_dataset["script_chunks"].append(src_chunk)
+        processed_dataset["script_summaries"].append(revised_summary)
+
+        elem_dict = {
+            "chunk_index": len(processed_dataset["script_chunks"]) - 1,
+            "chunk": src_chunk.strip(),
+            "summary": revised_summary.strip(),
+            "org_summary": original_summary.strip(),
+            "hallu_in_summary": list(hallu_summary_parts),
+            "feedbacks": list(feedbacks),
+        }
+        elem_dict_list.append(elem_dict)
+
+    agg_dict['script'] = summary['script']
+    agg_dict['summaries'] = ' '.join(processed_dataset["script_summaries"])
+
+    return {
+        "summary_sep": elem_dict_list,
+        "summary_agg": agg_dict,
+        "processed_dataset": processed_dataset
+    }
+
+def refine_summary_and_return_diff(summary, fact_score, api_key, model_id):
+    refined_summary = refine_summary(summary['processed_dataset'], fact_score, api_key, model_id)
+    diff = HtmlDiff().make_file(
+        summary['summary_agg']['summaries'].splitlines(),
+        refined_summary['summary_agg']['summaries'].splitlines(),
+        context=True
+    )
+    return diff
+
+def open_kg(kg_data):
+    if kg_data is None:
+        return "Please build the knowledge graph first."
+    try:
+        with open('refined_kg.html', 'r', encoding='utf-8') as f:
+            html_content = f.read()
+        return f'''
+            <iframe 
+                srcdoc="{html_content.replace('"', '&quot;')}"
+                style="width: 100%; height: 500px; border: none;"
+            ></iframe>
+        '''
+    except Exception as e:
+        return f'<div style="color: red;">Error reading KG file: {str(e)}</div>'
+
+def format_fact_score_output(fact_score_result):
+    if not fact_score_result:
+        return "No factuality analysis available"
+    
+    formatted_output = []
+    
+    # Overall score
+    formatted_output.append(f"Overall Factuality Score: {fact_score_result['fact_score']*100:.1f}%\n")
+    
+    # Individual chunk analysis
+    for i, chunk in enumerate(fact_score_result['summary_feedback_pairs'], 1):
+        formatted_output.append(f"\nChunk {i} Analysis:")
+        formatted_output.append("Original Text:")
+        formatted_output.append(f"{' '.join(chunk['summary_chunks'])}\n")
+        
+        if chunk['feedbacks']:
+            formatted_output.append("Feedback:")
+            feedbacks = [f"• {feedback}" for feedback in chunk['feedbacks'] if feedback.strip()]
+            formatted_output.extend(feedbacks)
+        
+        formatted_output.append("-" * 80)
+    
+    return "\n".join(formatted_output)
+
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # NarrativeFactScore: Script Factuality Evaluation
+        Evaluate and refine script summaries using narrative factuality scoring.
+        """
+    )
+    
+    with gr.Accordion("Model Settings", open=True):
+        with gr.Row():
+            api_key_input = gr.Textbox(
+                label="GPT API Key",
+                placeholder="Enter your GPT API key",
+                type="password",
+                scale=2
+            )
+            model_selector = gr.Dropdown(
+                choices=[
+                    "gpt-4o-mini",
+                    "gpt-4o",
+                    "gpt-4-turbo",
+                    "gpt-3.5-turbo-0125"
+                ],
+                value="gpt-4o",
+                label="Model Selection",
+                scale=1
+            )
+    
+    with gr.Tabs():
+        with gr.TabItem("Dataset Selection"):
+            with gr.Row():
+                dataset_selector = gr.Radio(
+                    choices=["MovieSum", "MENSA", "Custom"],
+                    label="Dataset",
+                    info="Choose the dataset or input custom script"
+                )
+                data_type_selector = gr.Radio(
+                    choices=["train", "validation", "test"],
+                    label="Split Type",
+                    info="Select data split",
+                    visible=True
+                )
+            name_list = gr.Dropdown(
+                choices=[],
+                label="Select Script",
+                info="Choose a script to analyze",
+                visible=True
+            )
+            custom_input = gr.Textbox(
+                label="Custom Script Input",
+                info="Enter scenes as a JSON list: ['scene1', 'scene2', ...]",
+                lines=10,
+                visible=False
+            )
+            narrative_output = gr.Textbox(
+                label="Script Content",
+                interactive=False,
+                lines=10
+            )
+            
+        with gr.TabItem("Knowledge Graph"):
+            with gr.Row():
+                generate_kg_button = gr.Button(
+                    "Generate Knowledge Graph",
+                    variant="primary"
+                )
+                open_kg_button = gr.Button("View Graph")
+            kg_status = gr.Textbox(
+                label="Status",
+                interactive=False
+            )
+            kg_viewer = gr.HTML(label="Knowledge Graph Visualization")
+            
+        with gr.TabItem("Summary Generation"):
+            generate_summary_button = gr.Button(
+                "Generate Initial Summary",
+                variant="primary"
+            )
+            summary_output = gr.Textbox(
+                label="Generated Summary",
+                interactive=False,
+                lines=5
+            )
+            calculate_score_button = gr.Button("Calculate Factuality Score")
+            fact_score_display = gr.Textbox(
+                label="Factuality Analysis",
+                interactive=False,
+                lines=10
+            )
+                
+        with gr.TabItem("Summary Refinement"):
+            refine_button = gr.Button(
+                "Refine Summary",
+                variant="primary"
+            )
+            refined_output = gr.HTML(label="Refined Summary with Changes")
+
+    # Hidden states
+    work_index = gr.State()
+    data_state = gr.State()
+    kg_output = gr.State()
+    summary_state = gr.State()
+    fact_score_state = gr.State()
+
+    # Event handlers
+    dataset_selector.change(
+        fn=lambda x: gr.update(visible=x in ["MovieSum", "MENSA"]),
+        inputs=[dataset_selector],
+        outputs=data_type_selector
+    )
+    
+    dataset_selector.change(
+        fn=update_name_list_interface,
+        inputs=[dataset_selector, data_type_selector],
+        outputs=[name_list, custom_input, narrative_output]
+    )
+    
+    name_list.change(
+        fn=get_narrative_content_with_index,
+        inputs=[dataset_selector, data_type_selector, name_list],
+        outputs=[narrative_output, work_index, data_state]
+    )
+    
+    generate_kg_button.click(
+        fn=build_kg_with_data,
+        inputs=[
+            data_state,      # data
+            work_index,      # work_index
+            custom_input,    # custom_scenes
+            api_key_input,   # api_key
+            model_selector   # model_id
+        ],
+        outputs=[kg_output, kg_status]
+    )
+    
+    open_kg_button.click(
+        fn=open_kg,
+        inputs=[kg_output],
+        outputs=kg_viewer
+    )
+    
+    generate_summary_button.click(
+        fn=generate_summary_with_data,
+        inputs=[data_state, work_index, custom_input, api_key_input, model_selector],
+        outputs=[summary_state, summary_output]
+    )
+    
+    calculate_score_button.click(
+        fn=lambda summary, kg, api_key, model: (
+            *calculate_narrative_fact_score(summary, kg, api_key, model),
+            format_fact_score_output(calculate_narrative_fact_score(summary, kg, api_key, model)[0])
+        ),
+        inputs=[summary_state, kg_output, api_key_input, model_selector],
+        outputs=[fact_score_state, fact_score_display]
+    )
+        
+    refine_button.click(
+        fn=refine_summary_and_return_diff,
+        inputs=[summary_state, fact_score_state, api_key_input, model_selector],
+        outputs=refined_output
+    )
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,200 @@
+absl-py==2.1.0
+accelerate==1.1.1
+aiofiles==23.2.1
+aiohappyeyeballs==2.4.4
+aiohttp==3.11.9
+aiosignal==1.3.1
+airportsdata==20241001
+alembic==1.14.0
+annotated-types==0.7.0
+anyio==4.6.2.post1
+asttokens
+async-timeout==5.0.1
+attrs==24.2.0
+banal==1.0.6
+bert-score==0.3.13
+bitsandbytes==0.44.1
+blis==1.0.1
+breadability==0.1.20
+catalogue==2.0.10
+certifi==2024.8.30
+chardet==5.2.0
+charset-normalizer==3.4.0
+click==8.1.7
+cloudpathlib==0.20.0
+cloudpickle==3.1.0
+cmake==3.31.1
+confection==0.1.5
+contourpy==1.3.1
+cycler==0.12.1
+cymem==2.0.10
+dataset==1.6.2
+datasets==2.16.0
+debugpy
+decorator
+dill==0.3.6
+diskcache==5.6.3
+distro==1.9.0
+docopt==0.6.2
+entrypoints
+evaluate==0.4.3
+exceptiongroup 
+executing
+fastapi==0.115.6
+ffmpy==0.4.0
+filelock==3.16.1
+FlagEmbedding==1.2.11
+fonttools==4.55.1
+frozenlist==1.5.0
+fsspec==2023.10.0
+gradio==5.8.0
+gradio_client==1.5.1
+greenlet==3.1.1
+h11==0.14.0
+httpcore==1.0.7
+httptools==0.6.4
+httpx==0.28.0
+huggingface-hub==0.26.3
+idna==3.10
+interegular==0.3.3
+ipdb==0.13.13
+jedi
+Jinja2==3.1.4
+jiter==0.8.0
+joblib==1.4.2
+jsonpickle==4.0.0
+jsonschema==4.23.0
+jsonschema-specifications==2024.10.1
+kiwisolver==1.4.7
+langcodes==3.5.0
+language_data==1.3.0
+lark==1.2.2
+lm-format-enforcer==0.10.1
+longdocfactscore==1.0.0
+lxml==5.3.0
+Mako==1.3.8
+marisa-trie==1.2.1
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matching==1.4
+matplotlib==3.9.3
+mdurl==0.1.2
+mpmath==1.3.0
+msgpack==1.1.0
+multidict==6.1.0
+multiprocess==0.70.14
+murmurhash==1.0.11
+nest_asyncio
+networkx==3.4.2
+ninja==1.11.1.2
+nltk==3.6.2
+numpy==2.0.2
+nvidia-cublas-cu11==11.11.3.6
+nvidia-cuda-cupti-cu11==11.8.87
+nvidia-cuda-nvrtc-cu11==11.8.89
+nvidia-cuda-runtime-cu11==11.8.89
+nvidia-cudnn-cu11==8.7.0.84
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.3.0.86
+nvidia-cusolver-cu11==11.4.1.48
+nvidia-cusparse-cu11==11.7.5.86
+nvidia-ml-py==12.560.30
+nvidia-nccl-cu11==2.20.5
+nvidia-nvtx-cu11==11.8.86
+openai==0.28.1
+orjson==3.10.12
+outlines==0.1.7
+outlines_core==0.1.17
+packaging
+pandas==2.2.3
+parso
+peft==0.13.2
+pexpect
+pickleshare
+pillow==11.0.0
+platformdirs
+preshed==3.0.9
+prometheus-fastapi-instrumentator==7.0.0
+prometheus_client==0.21.1
+prompt_toolkit
+propcache==0.2.1
+protobuf==5.29.0
+psutil
+ptyprocess
+pure_eval
+py-cpuinfo==9.0.0
+pyarrow==18.1.0
+pyarrow-hotfix==0.6
+pycountry==24.6.1
+pydantic==2.10.3
+pydantic_core==2.27.1
+pydub==0.25.1
+Pygments
+pyparsing==3.2.0
+pysbd==0.3.4
+python-dateutil
+python-dotenv==1.0.1
+python-multipart==0.0.19
+pytz==2024.2
+pyvis==0.3.2
+PyYAML==6.0.2
+pyzmq
+ray==2.40.0
+referencing==0.35.1
+regex==2024.11.6
+requests==2.32.3
+responses==0.18.0
+rich==13.9.4
+rouge_score==0.1.2
+rpds-py==0.22.1
+ruff==0.8.3
+safehttpx==0.1.6
+safetensors==0.4.5
+scikit-learn==1.5.2
+scipy==1.14.1
+seaborn==0.13.2
+semantic-version==2.10.0
+sentence-transformers==3.3.1
+sentencepiece==0.2.0
+shellingham==1.5.4
+six==1.16.0
+smart-open==7.0.5
+sniffio==1.3.1
+spacy==3.8.2
+spacy-legacy==3.0.12
+spacy-loggers==1.0.5
+SQLAlchemy==1.4.54
+srsly==2.4.8
+stack_data
+starlette==0.41.3
+sumy==0.11.0
+sympy==1.13.3
+tenacity==9.0.0
+thinc==8.3.2
+threadpoolctl==3.5.0
+tiktoken==0.8.0
+tokenizers==0.20.3
+tomli==2.2.1
+tomlkit==0.13.2
+torch==2.3.0
+tornado
+tqdm==4.67.1
+traitlets
+transformers==4.46.3
+triton==2.3.0
+typer==0.15.0
+typing_extensions
+tzdata==2024.2
+Unidecode==1.2.0
+urllib3==2.2.3
+uvicorn==0.32.1
+uvloop==0.21.0
+wasabi==1.1.3
+watchfiles==1.0.0
+wcwidth
+weasel==0.4.1
+websockets==14.1
+wrapt==1.17.0
+xformers==0.0.26.post1
+xxhash==3.5.0
+yarl==1.18.3

src/fact/narrativefactscore.py

@@ -0,0 +1,230 @@
+# Suppress annoying warnings from this issue which cannot be solved: https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md and transformers packages
+import warnings
+warnings.filterwarnings("ignore")
+
+import re
+import torch
+import torch.nn as nn
+import traceback
+from transformers import BartTokenizer, BartForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import numpy as np
+from nltk import sent_tokenize
+import logging
+import openai
+from tqdm import tqdm
+from sentence_transformers import SentenceTransformer, util
+from openai.error import (APIError, RateLimitError, ServiceUnavailableError,
+                          Timeout, APIConnectionError, InvalidRequestError)
+from tenacity import (before_sleep_log, retry, retry_if_exception_type,
+                      stop_after_delay, wait_random_exponential, stop_after_attempt)
+from .utils import break_down2scenes
+from .prompt import build_fact_prompt
+from .openai_api import openai_api_response
+
+
+logger = logging.getLogger(__name__)
+
+class OpenAIEmbedding:
+    def __init__(self, api_key, model="text-embedding-3-large"):
+        self.api_key = api_key
+        self.model = model
+        openai.api_key = api_key
+
+    @retry(retry=retry_if_exception_type((APIError, Timeout, RateLimitError,
+                                      ServiceUnavailableError, APIConnectionError)),
+           wait=wait_random_exponential(max=60), stop=stop_after_attempt(10),
+           before_sleep=before_sleep_log(logger, logging.WARNING))
+    def encode(self, texts, **kwargs):
+        if isinstance(texts, str):
+            texts = [texts]
+        
+        try:
+            response = openai.Embedding.create(
+                model=self.model,
+                input=texts,
+            )
+            
+            # Extract embeddings from response
+            embeddings = [item["embedding"] for item in response["data"]]
+            return np.array(embeddings)
+        
+        except Exception as e:
+            logger.error(f"Embedding API failed: {str(e)}")
+            return None
+
+class NarrativeFactScore():
+    def __init__(self, model="gpt-4o-mini", split_type="fast", checkpoint=None, api_key=None, model_id="gpt-4"):
+        self.sent_model = OpenAIEmbedding(api_key=api_key)
+        self.split_type = split_type
+        self.checkpoint = checkpoint
+        self.api_key = api_key
+        self.model_id = model_id
+        openai.api_key = api_key
+        
+        if model == "gptscore":
+            self.metric = GPTScore(model=self.model_id, api_key=self.api_key)
+            self.metric_function = self.metric.gpt_score
+        else:
+            raise ValueError("NarrativeFactScore currently only supports GPTScore")
+
+    def get_surrounding_sentences(self, sentence_array, ii):
+        if ii > 0 and ii < len(sentence_array) - 1:
+            sents = " ".join(np.array(sentence_array)[ii - 1 : ii + 1])
+        elif ii == 0:
+            sents = " ".join(np.array(sentence_array)[:2])
+        elif ii == len(sentence_array) - 1:
+            sents = " ".join(np.array(sentence_array)[ii - 1 :])
+        return sents
+
+    def group_into_sections(self, sentence_array, num_sent):
+        sectioned_sents = []
+        for ii in range(0, len(sentence_array), num_sent):
+            sectioned_sents.append(" ".join(sentence_array)[ii : ii + num_sent])
+        return sectioned_sents
+    
+    def split_sent(self, text):
+        text_list = []
+        if self.split_type == "fast":
+            for t in text.split('.'):
+                if len(t) == 0:
+                    continue
+                text_list.append(t)
+            return text_list
+        elif self.split_type == "fast_comma":
+            for t in re.split(r'[.,]', text):
+                if len(t) == 0:
+                    continue
+                text_list.append(t)
+            return text_list
+        elif self.split_type == "gpt":
+            prompt = build_fact_prompt(
+                prompt_template = './templates/atomic_fact.txt',
+                input_text_list=[text],
+            )
+            response = openai_api_response(prompt, model=self.model_id, api_key=self.api_key)
+            text_list = []
+            for res in response.split('\n'):
+                text_list.append(res.strip())
+            return text_list
+        else:
+            return None
+
+    def score_src_hyp_long(self, srcs, hyps, kgs):
+        all_scores = []
+        all_scores_per_sent = []
+        all_relevant_scenes = []
+        all_summary_chunks = []
+        all_feedback_list = []
+        # src is a list containing source documents.
+        # hyps is a list containing predicted documents
+        total_score = 0
+        for global_idx, (src, hyp) in enumerate(zip(tqdm(srcs), hyps)):
+            src_sents = break_down2scenes(src)
+            # Get embeddings using OpenAI API
+            sentence_embeddings_src = self.sent_model.encode(src_sents)
+            sentence_embeddings_kg = self.sent_model.encode(kgs)
+            
+            doc_scores = []
+            relevant_scenes = []
+            feedbacks = []
+            hyp_array = self.split_sent(hyp)
+            for idx, hyp_sentence in enumerate(hyp_array):
+                # Get embedding for hypothesis sentence
+                sentence_embeddings_hyp = self.sent_model.encode(hyp_sentence)
+                
+                # Calculate cosine similarity
+                scores = util.cos_sim(sentence_embeddings_hyp, sentence_embeddings_src)[0]
+                scores_kg = util.cos_sim(sentence_embeddings_hyp, sentence_embeddings_kg)[0]
+                
+                sorted_idxs = np.argsort(-1 * scores) # descending order
+                sorted_idxs_kg = np.argsort(-1 * scores_kg) # descending order
+                similar_src_sentences = []
+                similar_src_sentences_kg = []
+                triple = ''
+                
+                for sorted_idx, ii in enumerate(sorted_idxs_kg[0:1]):
+                    if sorted_idx == 0:
+                        triple += f'{kgs[ii]}'
+                    else:
+                        triple += f', {kgs[ii]}'
+                for ii in sorted_idxs[0:1]:
+                    similar_sents = src_sents[ii]
+                    similar_src_sentences.append(similar_sents)
+                
+                scores, feedback_list = self.metric_function(similar_src_sentences, [hyp_sentence for i in range(0, len(similar_src_sentences))], triple)
+                score = np.max(scores)
+                max_scene_idx = np.argmax(scores)
+                max_scene = similar_src_sentences[max_scene_idx]
+                feedback = feedback_list[max_scene_idx]
+                
+                doc_scores.append(int(score))
+                relevant_scenes.append(max_scene)
+                feedbacks.append(feedback)
+
+            doc_score = np.mean(doc_scores)
+            all_scores_per_sent.append(doc_scores)
+            all_scores.append(doc_score)
+            all_relevant_scenes.append(relevant_scenes)
+            all_summary_chunks.append(hyp_array)
+            all_feedback_list.append(feedbacks)
+            total_score += doc_score
+            if global_idx % 100 == 99:
+                print(f"Document mean {global_idx+1} Score: {total_score/(global_idx+1)} Score")
+        return all_scores, all_scores_per_sent, all_relevant_scenes, all_summary_chunks, all_feedback_list
+
+class GPTScore():
+    def __init__(self, model="gpt-4o", api_key=None, prompt='./templates/fact_score_kg.txt'):
+        self.max_length = 1024
+        self.model = model
+        self.api_key = api_key
+        self.prompt = prompt
+        openai.api_key = api_key
+    
+    @retry(retry=retry_if_exception_type((APIError, Timeout, RateLimitError,
+                                        ServiceUnavailableError, APIConnectionError, InvalidRequestError)),
+        wait=wait_random_exponential(max=60), stop=stop_after_attempt(10),
+        before_sleep=before_sleep_log(logger, logging.WARNING))
+    def gpt_inference(self, prompt):
+        prompt_messages = [{"role": "user", "content": prompt}]
+        try:
+            response = openai.ChatCompletion.create(
+                model=self.model,
+                messages=prompt_messages,
+                temperature=0,
+                api_key=self.api_key
+            )
+            response = response.choices[0].message.content
+        except InvalidRequestError:
+            response = 1
+        return response
+
+    def gpt_score(self, srcs, tgts, kgs, batch_size=4):
+        score_list = []
+        feedback_list = []
+
+        for i in range(len(srcs)):
+            src = srcs[i]
+            tgt = tgts[i]
+
+            prompt = build_fact_prompt(
+                prompt_template=self.prompt,
+                input_text_list=[src, kgs, tgt],
+            )
+            
+            try:
+                score = self.gpt_inference(prompt)
+                if '1' in score:
+                    score_list.append(float(1))
+                    feedback_list.append('')
+                else:
+                    score_list.append(float(0))
+                    feedback_list.append(score)
+            
+            except RuntimeError:
+                traceback.print_exc()
+                print(f"source: {src_list}")
+                print(f"target: {tgt_list}")
+                exit(0)
+                
+        return score_list, feedback_list

src/fact/openai_api.py

@@ -0,0 +1,52 @@
+import json
+import logging
+import os
+from pathlib import Path
+
+import openai
+from dotenv import load_dotenv
+from openai.error import (APIError, RateLimitError, ServiceUnavailableError,
+                          Timeout, APIConnectionError, InvalidRequestError)
+from tenacity import (before_sleep_log, retry, retry_if_exception_type,
+                      stop_after_delay, wait_random_exponential, stop_after_attempt)
+from tiktoken import Encoding, encoding_for_model
+
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+# This value is set by OpenAI for the selected model and cannot be changed.
+MAX_MODEL_TOKEN_COUNT = 4096
+# This value can be changed.
+MAX_RESPONSE_TOKEN_COUNT = 512
+
+
+
+@retry(retry=retry_if_exception_type((APIError, Timeout, RateLimitError,
+                                      ServiceUnavailableError, APIConnectionError, InvalidRequestError)),
+       wait=wait_random_exponential(max=60), stop=stop_after_attempt(10),
+       before_sleep=before_sleep_log(logger, logging.WARNING))
+def openai_api_response(prompt, model, api_key, save_path=None):
+    """
+    Use a prompt to make a request to the OpenAI API and save the response to a
+    JSON file.
+    """
+    openai.api_key = api_key
+    try:
+        prompt_messages = [{"role": "user", "content": prompt}]
+        response = openai.ChatCompletion.create(
+            model=model, messages=prompt_messages, temperature=0)
+        finish_reason = response.choices[0].finish_reason
+        if finish_reason != 'stop':
+            logger.error(f'`finish_reason` is `{finish_reason}` for {save_path}.')
+        save_data = {'model': response.model, 'usage': response.usage,
+                     'finish_reason': finish_reason,
+                     'prompt_messages': prompt_messages,
+                     'response': response.choices[0].message.content}
+    except InvalidRequestError:
+        logger.error(f'InvalidRequestError encountered 10 times. Returning empty string for {save_path}.')
+        save_data = {'model': None, 'usage': None,
+                     'finish_reason': 'invalid_request',
+                     'prompt_messages': prompt_messages,
+                     'response': ' '}
+    return save_data['response']

src/fact/prompt.py

@@ -0,0 +1,28 @@
+import os
+from typing import List, Optional
+
+
+def build_fact_prompt(
+    prompt_template:str,
+    input_text_list:List[str],
+    chat_mode:Optional[str] = None) -> str:
+    
+    """_summary_
+    chat_mode(str) : 'hf-chat', 'kullm', 'None'
+
+    Returns:
+        _type_: _description_
+    """
+    
+    if os.path.isfile(prompt_template):
+        with open(prompt_template,'r') as f:
+            prompt_template = f.read() 
+    else:
+        pass
+
+    assert isinstance(prompt_template, str)
+    
+    prompt = prompt_template.format(*input_text_list)
+    
+    
+    return prompt

src/fact/utils.py

@@ -0,0 +1,47 @@
+import re
+
+def delete_special(pre_text, character_list):
+    for c in character_list:
+        pre_text = pre_text.replace(c, "")
+    return pre_text
+
+def break_down2scenes(text: str):
+    # Split the text based on the 's#' pattern
+    scenes = re.split(r'(s#\d+)', text)
+
+    # Remove empty elements from the split results
+    scenes = [scene for scene in scenes if scene.strip()]
+
+    scenes_list = []
+    current_scene_number = None
+
+    for i in range(0, len(scenes), 2):  # Process the 's#' marker and its corresponding text as pairs
+        scene_marker = scenes[i].strip()
+        try:
+            scene_number = int(scene_marker.split('#')[1])  # Extract the number part
+        except:
+            if len(scenes) % 2 == 1:
+                return [scenes[0]]
+            import ipdb;ipdb.set_trace(context=10)
+        scene_text = scenes[i+1].strip() if i+1 < len(scenes) else ""
+
+        # Check if the scene numbers are in the correct sequence
+        if current_scene_number is not None:
+            expected_scene_number = current_scene_number + 1
+            if scene_number != expected_scene_number:
+                raise ValueError(f"Unexpected scene number: {scene_number}, expected {expected_scene_number}")
+
+        # Store the scene number and its corresponding text together
+        scenes_list.append({
+            'detected_scene_number': scene_number,
+            'text': f"{scene_marker}\n{scene_text}".strip()
+        })
+
+    filtered_scene_list = []
+    scene_number = 0
+    for scene_dict in scenes_list:
+        detected_scene_number = int(scene_dict['detected_scene_number'])
+        filtered_scene_list.append(scene_dict['text'])
+        scene_number = detected_scene_number
+
+    return filtered_scene_list

src/kg/generate_kg.py

@@ -0,0 +1,253 @@
+# src.kg.generate_kg.py
+import pickle
+from collections import defaultdict, Counter
+from contextlib import redirect_stdout
+from pathlib import Path
+import json
+import argparse
+import os
+import openai
+import time
+import numpy as np
+
+import networkx as nx
+from pyvis.network import Network
+from tqdm import tqdm
+from contextlib import redirect_stdout
+
+
+
+from .knowledge_graph import generate_knowledge_graph
+from .openai_api import load_response_text
+from .save_triples import get_response_save_path
+from .utils import set_up_logging
+
+logger = set_up_logging('generate-knowledge-graphs-books.log')
+KNOWLEDGE_GRAPHS_DIRECTORY_PATH = Path('../knowledge-graphs_new')
+
+
+"""def gpt_inference(system_instruction, prompt, retries=10, delay=5):
+    # api 
+    messages = [{"role": "system", "content": system_instruction}, 
+                {"role": "user", "content": prompt}]
+    
+    for attempt in range(retries):
+        try:
+            response = openai.ChatCompletion.create(
+                model='gpt-4o-mini-2024-07-18',
+                messages=messages,
+                temperature=0.0,
+                max_tokens=128,
+                top_p=0.5,
+                frequency_penalty=0,
+                presence_penalty=0
+            )
+            result = response['choices'][0]['message']['content']
+            return result
+        except openai.error.APIError as e:
+            
+            time.sleep(delay)
+            continue"""
+
+
+def generate_knowledge_graph_for_scripts(book, idx, save_path):
+    """
+    Use the responses from the OpenAI API to generate a knowledge graph for a
+    book.
+    """
+    response_texts = defaultdict(list)
+    project_gutenberg_id = book['id']
+    for chapter in book['chapters']:
+        chapter_index = chapter['index']
+        chapter_responses_directory = get_response_save_path(
+            idx, save_path, project_gutenberg_id, chapter_index)
+        for response_path in chapter_responses_directory.glob('*.json'):
+            response_text = load_response_text(response_path)
+            response_texts[chapter_index].append(response_text)
+    knowledge_graph = generate_knowledge_graph(response_texts, project_gutenberg_id)
+    return knowledge_graph
+
+def generate_knowledge_graph_for_scripts(book, idx, response_list):
+    """
+    Use the responses from the OpenAI API to generate a knowledge graph for a
+    book.
+    """
+
+    response_texts = defaultdict(list)
+    project_gutenberg_id = book['id']
+    for chapter in book['chapters']:
+        chapter_index = chapter['index']
+        for response in response_list:
+            response_texts[chapter_index].append(response['response'])
+    knowledge_graph = generate_knowledge_graph(response_texts, project_gutenberg_id)
+    return knowledge_graph
+
+
+def save_knowledge_graph(knowledge_graph,
+                         project_gutenberg_id, save_path):
+    """Save a knowledge graph to a `pickle` file."""
+    save_path = save_path / 'kg.pkl'
+    save_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(save_path, 'wb') as knowledge_graph_file:
+        pickle.dump(knowledge_graph, knowledge_graph_file)
+
+
+def load_knowledge_graph(project_gutenberg_id, save_path):
+    """Load a knowledge graph from a `pickle` file."""
+    save_path = save_path / 'kg.pkl'
+    with open(save_path, 'rb') as knowledge_graph_file:
+        knowledge_graph = pickle.load(knowledge_graph_file)
+    return knowledge_graph
+
+
+def display_knowledge_graph(knowledge_graph, save_path):
+    """Display a knowledge graph using pyvis."""
+    # Convert the knowledge graph into a format that can be displayed by pyvis.
+    # Merge all edges with the same subject and object into a single edge.
+    pyvis_graph = nx.MultiDiGraph()
+    for node in knowledge_graph.nodes:
+        pyvis_graph.add_node(str(node), label='\n'.join(node.names),
+                             shape='box')
+    for edge in knowledge_graph.edges(data=True):
+        subject = str(edge[0])
+        object_ = str(edge[1])
+        predicate = edge[2]['predicate']
+        chapter_index = edge[2]['chapter_index']
+        if pyvis_graph.has_edge(subject, object_):
+            pyvis_graph[subject][object_][0].update(
+                title=(f'{pyvis_graph[subject][object_][0]["title"]}\n'
+                       f'{predicate}')) # f'{predicate} ({chapter_index})'))
+        else:
+            pyvis_graph.add_edge(subject, object_,
+                                 title=f'{predicate}') # title=f'{predicate} ({chapter_index})')
+    network = Network(height='99vh', directed=True, bgcolor='#262626',
+                      cdn_resources='remote')
+    network.set_options('''
+    const options = {
+        "interaction": {
+            "tooltipDelay": 0
+        },
+        "physics": {
+            "forceAtlas2Based": {
+                "gravitationalConstant": -50,
+                "centralGravity": 0.01,
+                "springLength": 100,
+                "springConstant": 0.08,
+                "damping": 0.4,
+                "avoidOverlap": 0
+            },
+            "solver": "forceAtlas2Based"
+        }
+    }''')
+    network.from_nx(pyvis_graph)
+    save_path.parent.mkdir(parents=True, exist_ok=True)
+    # `show()` tries to print the name of the HTML file to the console, so
+    # suppress it.
+    with redirect_stdout(None):
+        network.show(str(save_path), notebook=False)
+    logger.info(f'Saved pyvis knowledge graph to {save_path}.')
+
+def fuse_subject(subjects):
+    subject_list = subjects.split('/')
+    if len(subject_list) == 1:
+        return subject_list[0]
+    flag = 0
+    striped_subject_list = []
+    len_list = []
+    for subject in subject_list:
+        striped_subject_list.append(subject.strip())
+        len_list.append(len(subject))
+    idx = np.argmin(len_list)
+    for subject in striped_subject_list:
+        if striped_subject_list[idx] in subject:
+            flag += 1
+
+    if flag == len(striped_subject_list):
+        return striped_subject_list[idx]
+    else:
+        return subjects
+
+def init_kg(script, idx, response_list):
+    """
+    Generate knowledge graphs for book in the books dataset using saved
+    responses from the OpenAI API.
+    """
+    knowledge_graph = generate_knowledge_graph_for_scripts(script, idx, response_list)
+    return knowledge_graph
+
+def refine_kg(knowledge_graph, idx, topk):
+    result = []
+    edge_count = Counter()
+    for edge in knowledge_graph.edges(data=True):
+        subject = str(edge[0])
+        object_ = str(edge[1])
+        edge_count[subject] += 1
+        edge_count[object_] += 1
+
+    # 엣지가 많은 상위 k개의 노드 선택
+    top_k_nodes = [node for node, count in edge_count.most_common(topk)]
+
+    # 상위 k개 노드 간의 모든 관계를 수집
+    rel_dict = {}
+    for edge in knowledge_graph.edges(data=True):
+        subject = str(edge[0])
+        object_ = str(edge[1])
+        if subject in top_k_nodes and object_ in top_k_nodes:
+            predicate = edge[2]['predicate']
+            chapter_index = edge[2]['chapter_index']
+            count = edge[2]['count']
+            key = f"{subject}\t{object_}"
+            if key not in rel_dict:
+                rel_dict[key] = []
+            rel_dict[key].append((predicate, chapter_index, count))
+
+    # 시각화 코드
+    pyvis_graph = nx.MultiDiGraph()
+    for node in top_k_nodes:
+        pyvis_graph.add_node(node, label=node, shape='box')
+
+    for key, relations in rel_dict.items():
+        subject, object_ = key.split('\t')
+        for relation in relations:
+            predicate, chapter_index, count = relation
+            if 'output' in predicate:
+                continue
+            if count >= 2:
+                if pyvis_graph.has_edge(subject, object_):
+                    pyvis_graph[subject][object_][0]['title'] += f', {predicate}'
+                else:
+                    pyvis_graph.add_edge(subject, object_, title=f'{predicate}')
+
+    network = Network(height='99vh', directed=True, bgcolor='#262626', cdn_resources='remote')
+    network.from_nx(pyvis_graph)
+
+    with redirect_stdout(None):
+        network.show('refined_kg.html', notebook=False)
+
+    for key, relations in rel_dict.items():
+        subject, object_ = key.split('\t')
+        
+        for relation in relations:
+            predicate, chapter_index, count = relation
+            
+            if 'output' in predicate:
+                continue
+
+            subject = fuse_subject(subject)
+            object_ = fuse_subject(object_)
+
+            relationship = {
+                'subject': subject,
+                'predicate': predicate,
+                'object': object_,
+                'chapter_index': chapter_index,
+                'count': count,
+                'subject_node_count': edge_count[subject],
+                'object_node_count': edge_count[object_] 
+            }
+
+            if count >= 2:
+                result.append(relationship)
+    
+    return result
+    

src/kg/knowledge_graph.py

@@ -0,0 +1,301 @@
+# src.kg.knowledge_graph.py
+import itertools
+import logging
+import re
+from collections import defaultdict
+from itertools import combinations, product
+from pathlib import Path
+
+import networkx as nx
+
+from .utils import strip_and_remove_empty_strings
+
+logger = logging.getLogger(__name__)
+
+PROMPT_FILE_PATH = Path('templates/story-prompt.txt')
+MAX_RESPONSE_EDGE_COUNT = 15
+MAX_PREDICATE_WORD_COUNT = 5
+MAX_POSSESSION_WORD_COUNT = 2
+MAX_MERGEABLE_NODE_EDGE_COUNT = 2
+MIN_NODE_EDGE_COUNT = 1
+
+
+class NamedEntity:
+    """A knowledge graph node representing a named entity."""
+
+    def __init__(self, names):
+        self.names = names
+
+    def __repr__(self):
+        return ' / '.join(self.names)
+
+def remove_number_prefix(text):
+
+    clean_text = re.sub(r'^\d+\.\s*', '', text)
+    return clean_text
+
+def parse_response_text(response_text, identifier, are_edges_numbered=True):
+    """
+    Parse a response text from the OpenAI model into names (a list of names for
+    each entity) and edges (relations between entities). `identifier` is a
+    string used to identify the response text in error messages.
+    """
+
+    lines = strip_and_remove_empty_strings(response_text.split('\n'))
+
+    if 'Named entities' not in lines[0]:
+        logger.error(f'{identifier}: First line of response text does not '
+                     f'start with "Named entities:". ("{lines[0]}")')
+        return [], []
+    mode = 'names'
+    names = []
+    edges = []
+    for line in lines[1:]:
+        if 'Knowledge graph edges' in line:
+            mode = 'edges'
+            continue
+        if mode == 'names':
+            if line.startswith('-'):
+                line = line[1:]
+            name_group = strip_and_remove_empty_strings(line.split(' / '))
+            name_group = [remove_number_prefix(name) for name in name_group]
+            names.append(name_group)
+        elif mode == 'edges':
+            if are_edges_numbered:
+                if not re.match(r'^\d{1,2}\. ', line):
+                    break
+                if int(line.split('.')[0]) > MAX_RESPONSE_EDGE_COUNT:
+                    break;
+                line = line[3:]
+            edge_components = strip_and_remove_empty_strings(line.split(';'))
+            if len(edge_components) not in (2, 3):
+                continue
+            subjects = strip_and_remove_empty_strings(
+                edge_components[0].split(','))
+            predicate = edge_components[1]
+            if len(edge_components) == 3:
+                objects = strip_and_remove_empty_strings(
+                    edge_components[2].split(','))
+            else:
+                objects = [None]
+            for subject, object_ in product(subjects, objects):
+                edge = (subject, predicate, object_)
+                edges.append(edge)
+    if not names:
+        logger.error(f'{identifier}: No names were parsed from the response '
+                     f'text.')
+    if not edges:
+        logger.error(f'{identifier}: No edges were parsed from the response '
+                     f'text.')
+
+    return names, edges
+
+
+def generate_names_graph(names):
+    """
+    Generate a graph of names where the nodes are names and the edges indicate
+    that two names refer to the same entity.
+    """
+    names_graph = nx.Graph()
+    for name_group in names:
+        for name in name_group:
+            names_graph.add_node(name)
+        for name_pair in combinations(name_group, 2):
+            names_graph.add_edge(*name_pair)
+    return names_graph
+
+
+def expand_contracted_possessive(predicate, names):
+    """
+    Check if a predicate is of the form "<owner>'s <possession>", where the
+    owner is a named entity. If so, return a predicate of the form
+    "<possession> of" and an object of the form "<owner>".
+    """
+    match = re.search(
+        fr'\'s\s\w+(?:\s\w+)'
+        fr'{{0,{MAX_POSSESSION_WORD_COUNT - 1}}}$', predicate)
+    if not match:
+        return predicate, None
+    apostrophe_index = match.start()
+    owner = next(
+        (name for name in names
+         if predicate[:apostrophe_index].endswith(name)), None)
+    if owner is None:
+        return predicate, None
+    possession = predicate[apostrophe_index + 2:].strip()
+    predicate = (f'{predicate[:apostrophe_index - len(owner)].strip()} '
+                 f'{possession} of')
+    object_ = owner
+    return predicate, object_
+
+
+def does_duplicate_edge_exist(knowledge_graph, subject, predicate, object_):
+    """
+    Check if an edge with a given subject, predicate, and object already exists
+    in a knowledge graph. If it exists, return the edge data; otherwise, return None.
+    """
+    for edge in knowledge_graph.edges(subject, data=True):
+        if edge[1] == object_ and edge[2]['predicate'] == predicate:
+            return edge
+    return None
+
+
+def add_edge_to_knowledge_graph(knowledge_graph, names, edge, max_predicate_word_count, **edge_attributes):
+    """Add an edge to a knowledge graph, updating count if the edge already exists."""
+    subject, predicate, object_ = edge
+    if subject not in names:
+        return
+    if object_ is not None and object_ not in names:
+        predicate += f' {object_}'
+        object_ = None
+    if object_ is None:
+        object_at_end_of_predicate = next(
+            (name for name in names if predicate.endswith(' ' + name)), None)
+        if object_at_end_of_predicate is not None:
+            object_ = object_at_end_of_predicate
+            predicate = predicate[:-len(object_)].strip()
+        else:
+            predicate, object_ = expand_contracted_possessive(predicate, names)
+    while predicate.endswith(('.', ',', '!', '?')):
+        predicate = predicate[:-1]
+    if (max_predicate_word_count and len(predicate.split()) > max_predicate_word_count):
+        return
+    if subject == object_:
+        return
+    if object_ is None:
+        object_ = subject
+    subject_node = next((node for node in knowledge_graph.nodes if subject in node.names), None)
+    object_node = next((node for node in knowledge_graph.nodes if object_ in node.names), None)
+
+    if subject_node is None or object_node is None:
+        return
+
+    existing_edge = does_duplicate_edge_exist(knowledge_graph, subject_node, predicate, object_node)
+    if existing_edge:
+        existing_edge[2]['count'] += 1
+    else:
+        knowledge_graph.add_edge(subject_node, object_node, predicate=predicate, count=1, **edge_attributes)
+
+
+def initialize_knowledge_graph(names_graph, edges):
+    """
+    Initialize a knowledge graph from a graph of names and a dictionary of
+    edges grouped by chapter index.
+    """
+    names = set(names_graph.nodes)
+    knowledge_graph = nx.MultiDiGraph()
+    for name in names:
+        knowledge_graph.add_node(NamedEntity({name}))
+    for chapter_index, chapter_edges in edges.items():
+        for edge in chapter_edges:
+            add_edge_to_knowledge_graph(
+                knowledge_graph, names, edge,
+                max_predicate_word_count=MAX_PREDICATE_WORD_COUNT,
+                chapter_index=chapter_index)
+    return knowledge_graph
+
+
+def get_node_edge_count(knowledge_graph, node):
+    """
+    Get the number of edges for a node in a knowledge graph, excluding
+    self-loops.
+    """
+    edges = (set(knowledge_graph.in_edges(node))
+             | set(knowledge_graph.out_edges(node)))
+    edge_count = sum(1 for edge in edges if edge[0] is not edge[1])
+    return edge_count
+
+
+def merge_nodes(knowledge_graph, nodes_to_merge):
+    """
+    Merge a list of nodes in a knowledge graph into one node, combining their
+    sets of names and preserving their edges.
+    """
+    merged_node = NamedEntity(set())
+    for node in nodes_to_merge:
+        merged_node.names.update(node.names)
+    knowledge_graph.add_node(merged_node)
+    for node in nodes_to_merge:
+        for edge in itertools.chain(knowledge_graph.out_edges(node, data=True),
+                                    knowledge_graph.in_edges(node, data=True)):
+            subject, object_, attributes = edge
+            if (does_duplicate_edge_exist(knowledge_graph, merged_node,
+                                          attributes['predicate'], object_)
+                    or does_duplicate_edge_exist(knowledge_graph, subject,
+                                                 attributes['predicate'],
+                                                 merged_node)):
+                continue
+            if subject is object_:
+                knowledge_graph.add_edge(merged_node, merged_node,
+                                         **attributes)
+            if subject is node:
+                knowledge_graph.add_edge(merged_node, object_, **attributes)
+            else:
+                knowledge_graph.add_edge(subject, merged_node, **attributes)
+        knowledge_graph.remove_node(node)
+
+def merge_same_entity_nodes(knowledge_graph, names_graph):
+    """
+    Using a graph of names, merge nodes in a knowledge graph corresponding to
+    the same entity.
+    """
+    for name_pair in names_graph.edges:
+        first_node = next((node for node in knowledge_graph.nodes
+                           if name_pair[0] in node.names), None)
+        if first_node is None:
+            continue
+        if name_pair[1] in first_node.names:
+            continue
+        second_node = next((node for node in knowledge_graph.nodes
+                            if name_pair[1] in node.names), None)
+        if second_node is None:
+            continue
+        if knowledge_graph.has_edge(first_node, second_node):
+            continue
+        first_node_edge_count = get_node_edge_count(knowledge_graph,
+                                                    first_node)
+        second_node_edge_count = get_node_edge_count(knowledge_graph,
+                                                     second_node)
+        if (first_node_edge_count > MAX_MERGEABLE_NODE_EDGE_COUNT
+                and second_node_edge_count > MAX_MERGEABLE_NODE_EDGE_COUNT):
+            continue
+        merge_nodes(knowledge_graph, [first_node, second_node])
+
+
+
+def remove_nodes_with_few_edges(knowledge_graph):
+    """
+    Remove nodes that have fewer than `MIN_NODE_EDGE_COUNT` edges (excluding
+    self-loops) from a knowledge graph. Repeat until no more nodes are removed.
+    """
+    while True:
+        nodes_to_remove = []
+        for node in knowledge_graph.nodes:
+            edge_count = get_node_edge_count(knowledge_graph, node)
+            if edge_count < MIN_NODE_EDGE_COUNT:
+                nodes_to_remove.append(node)
+        if not nodes_to_remove:
+            break
+        knowledge_graph.remove_nodes_from(nodes_to_remove)
+
+
+def generate_knowledge_graph(response_texts, project_gutenberg_index):
+    """
+    Use OpenAI API response texts grouped by chapter index to generate a
+    knowledge graph for a book.
+    """
+    names = []
+    edges = defaultdict(list)
+    for chapter_index, chapter_response_texts in response_texts.items():
+        for response_text in chapter_response_texts:
+            identifier = (f'Book {project_gutenberg_index}, chapter '
+                          f'{chapter_index}')
+            chapter_segment_names, chapter_segment_edges = parse_response_text(
+                response_text, identifier)
+            names.extend(chapter_segment_names)
+            edges[chapter_index].extend(chapter_segment_edges)
+    names_graph = generate_names_graph(names)
+    knowledge_graph = initialize_knowledge_graph(names_graph, edges)
+    merge_same_entity_nodes(knowledge_graph, names_graph)
+    remove_nodes_with_few_edges(knowledge_graph)
+    return knowledge_graph

src/kg/main.py

@@ -0,0 +1,23 @@
+# 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

src/kg/openai_api.py

@@ -0,0 +1,101 @@
+# src.kg.openai_api.py
+import json
+import logging
+import os
+from pathlib import Path
+
+import openai
+from dotenv import load_dotenv
+from openai.error import (APIError, RateLimitError, ServiceUnavailableError,
+                          Timeout, APIConnectionError, InvalidRequestError)
+from tenacity import (before_sleep_log, retry, retry_if_exception_type,
+                      stop_after_delay, wait_random_exponential, stop_after_attempt)
+from tiktoken import Encoding, encoding_for_model
+
+logger = logging.getLogger(__name__)
+
+load_dotenv()
+
+# This value is set by OpenAI for the selected model and cannot be changed.
+MAX_MODEL_TOKEN_COUNT = 4096
+# This value can be changed.
+MAX_RESPONSE_TOKEN_COUNT = 512
+RESPONSES_DIRECTORY_PATH = Path('../openai-api-responses-new')
+
+
+def get_openai_model_encoding(model_id):
+    """Get the encoding (tokenizer) for the OpenAI model."""
+    return encoding_for_model(model_id)
+
+
+def get_max_chapter_segment_token_count(prompt: str, model_id: str) -> int:
+    """
+    Calculate the maximum number of tokens that a chapter segment may contain
+    given the prompt.
+    """
+    encoding = get_openai_model_encoding(model_id)
+    # `encode_ordinary()` ignores special tokens and is slightly faster than
+    # `encode()`.
+    prompt_token_count = len(encoding.encode_ordinary(prompt))
+    # Subtract 8 for tokens added by OpenAI in the prompt and response (refer
+    # to https://platform.openai.com/docs/guides/chat/managing-tokens for
+    # details).
+    # Subtract 1 for the newline added below to the end of the prompt.
+    # This calculation does not have to be exact.
+    max_chapter_segment_token_count = (MAX_MODEL_TOKEN_COUNT
+                                       - MAX_RESPONSE_TOKEN_COUNT
+                                       - prompt_token_count - 8 - 1)
+    return max_chapter_segment_token_count
+
+
+@retry(retry=retry_if_exception_type((APIError, Timeout, RateLimitError,
+                                      ServiceUnavailableError, APIConnectionError, InvalidRequestError)),
+       wait=wait_random_exponential(max=60), stop=stop_after_attempt(10),
+       before_sleep=before_sleep_log(logger, logging.WARNING))
+def save_openai_api_response(prompt_messages):
+    """
+    Use a prompt to make a request to the OpenAI API and return the response data.
+    """
+    
+    openai.api_key = prompt_messages[0]['api_key']  # Set the API key for OpenAI
+    model_id = prompt_messages[0]['model_id']  # Get the model ID from the prompt messages
+    prompt_messages[0].pop('api_key')  # Remove the API key from the prompt messages 
+    prompt_messages[0].pop('model_id')  # Remove the model ID from the prompt messages
+    
+    try:
+        logger.info('Calling OpenAI API...')
+        response = openai.ChatCompletion.create(
+            model=model_id, messages=prompt_messages, temperature=0
+        )
+        finish_reason = response.choices[0].finish_reason
+        if finish_reason != 'stop':
+            logger.error(f'`finish_reason` is `{finish_reason}`.')
+        
+        save_data = {
+            'model': response.model,
+            'usage': response.usage,
+            'finish_reason': finish_reason,
+            'prompt_messages': prompt_messages,
+            'response': response.choices[0].message.content
+        }
+    except InvalidRequestError:
+        logger.error('InvalidRequestError encountered 10 times. Returning empty response.')
+        save_data = {
+            'model': None,
+            'usage': None,
+            'finish_reason': 'invalid_request',
+            'prompt_messages': prompt_messages,
+            'response': ' '
+        }
+
+    return save_data
+
+
+def load_response_text(save_path):
+    """
+    Load the response text from a JSON file containing response data from the
+    OpenAI API.
+    """
+    with open(save_path, 'r') as save_file:
+        save_data = json.load(save_file)
+    return save_data['response']

src/kg/preprocess.py

@@ -0,0 +1,20 @@
+
+
+
+def split_scenes(script):
+    # TBD
+    return script
+
+def preprocess(scene_list, idx):
+    script_dict = {}
+    script_dict['id'] = idx
+    script_dict['chapters'] = []
+
+    elem_dict = {}
+    elem_dict['index'] = 1
+    elem_dict['text'] = scene_list
+    elem_dict['summaries'] = ""
+
+    script_dict['chapters'].append(elem_dict)
+
+    return script_dict

src/kg/save_triples.py

@@ -0,0 +1,215 @@
+# src.kg.save_triples.py
+from pathlib import Path
+import json
+import argparse
+import os
+
+from pysbd import Segmenter
+from tiktoken import Encoding
+
+from .knowledge_graph import PROMPT_FILE_PATH
+from .openai_api import (RESPONSES_DIRECTORY_PATH,
+                        get_max_chapter_segment_token_count,
+                        get_openai_model_encoding, save_openai_api_response)
+from .utils import (execute_function_in_parallel, set_up_logging,
+                   strip_and_remove_empty_strings)
+
+logger = set_up_logging('openai-api-scripts.log')
+
+
+def get_paragraphs(text):
+    """Split a text into paragraphs."""
+    paragraphs = strip_and_remove_empty_strings(text.split('\n\n'))
+    # Convert all whitespace into single spaces.
+    paragraphs = [' '.join(paragraph.split()) for paragraph in paragraphs]
+    return paragraphs
+
+
+def combine_text_subunits_into_segments(subunits, join_string,
+                                        encoding: Encoding,
+                                        max_token_count):
+    """
+    Combine subunits of text into segments that do not exceed a maximum number
+    of tokens.
+    """
+    # `encode_ordinary_batch()` ignores special tokens and is slightly faster
+    # than `encode_batch()`.
+    subunit_token_counts = [len(tokens) for tokens
+                            in encoding.encode_ordinary_batch(subunits)]
+    join_string_token_count = len(encoding.encode_ordinary(join_string))
+    total_token_count = (sum(subunit_token_counts) + join_string_token_count
+                         * (len(subunits) - 1))
+    if total_token_count <= max_token_count:
+        return [join_string.join(subunits)]
+    # Calculate the approximate number of segments and the approximate number
+    # of tokens per segment, in order to keep the segment lengths roughly
+    # equal.
+    approximate_segment_count = total_token_count // max_token_count + 1
+    approximate_segment_token_count = round(total_token_count
+                                            / approximate_segment_count)
+    segments = []
+    current_segment_subunits = []
+    current_segment_token_count = 0
+    for i, (subunit, subunit_token_count) in enumerate(
+            zip(subunits, subunit_token_counts)):
+        # The token count if the current subunit is added to the current
+        # segment.
+        extended_segment_token_count = (current_segment_token_count
+                                        + join_string_token_count
+                                        + subunit_token_count)
+        # Add the current subunit to the current segment if it results in a
+        # token count that is closer to the approximate segment token count
+        # than the current segment token count.
+        if (extended_segment_token_count <= max_token_count
+                and abs(extended_segment_token_count
+                        - approximate_segment_token_count)
+                <= abs(current_segment_token_count
+                       - approximate_segment_token_count)):
+            current_segment_subunits.append(subunit)
+            current_segment_token_count = extended_segment_token_count
+        else:
+            segment = join_string.join(current_segment_subunits)
+            segments.append(segment)
+            # If it is possible to join the remaining subunits into a single
+            # segment, do so. Additionally, add the current subunit as a
+            # segment if it is the last subunit.
+            if (sum(subunit_token_counts[i:]) + join_string_token_count
+                    * (len(subunits) - i - 1) <= max_token_count
+                    or i == len(subunits) - 1):
+                segment = join_string.join(subunits[i:])
+                segments.append(segment)
+                break
+            current_segment_subunits = [subunit]
+            current_segment_token_count = subunit_token_count
+    return segments
+
+
+def split_long_sentences(sentences, encoding: Encoding,
+                         max_token_count):
+    """
+    Given a list of sentences, split sentences that exceed a maximum number of
+    tokens into multiple segments.
+    """
+    token_counts = [len(tokens) for tokens
+                    in encoding.encode_ordinary_batch(sentences)]
+    split_sentences = []
+    for sentence, token_count in zip(sentences, token_counts):
+        if token_count > max_token_count:
+            words = sentence.split()
+            segments = combine_text_subunits_into_segments(
+                words, ' ', encoding, max_token_count)
+            split_sentences.extend(segments)
+        else:
+            split_sentences.append(sentence)
+    return split_sentences
+
+
+def split_long_paragraphs(paragraphs, encoding: Encoding,
+                          max_token_count):
+    """
+    Given a list of paragraphs, split paragraphs that exceed a maximum number
+    of tokens into multiple segments.
+    """
+    token_counts = [len(tokens) for tokens
+                    in encoding.encode_ordinary_batch(paragraphs)]
+    split_paragraphs = []
+    for paragraph, token_count in zip(paragraphs, token_counts):
+        if token_count > max_token_count:
+            sentences = Segmenter().segment(paragraph)
+            sentences = split_long_sentences(sentences, encoding,
+                                             max_token_count)
+            segments = combine_text_subunits_into_segments(
+                sentences, ' ', encoding, max_token_count)
+            split_paragraphs.extend(segments)
+        else:
+            split_paragraphs.append(paragraph)
+    return split_paragraphs
+
+
+def get_chapter_segments(chapter_text, encoding: Encoding,
+                         max_token_count):
+    """
+    Split a chapter text into segments that do not exceed a maximum number of
+    tokens.
+    """
+    paragraphs = get_paragraphs(chapter_text)
+    paragraphs = split_long_paragraphs(paragraphs, encoding, max_token_count)
+    chapter_segments = combine_text_subunits_into_segments(
+        paragraphs, '\n', encoding, max_token_count)
+    return chapter_segments
+
+
+def get_response_save_path(idx, save_path, project_gutenberg_id,
+                           chapter_index = None,
+                           chapter_segment_index = None,
+                           chapter_segment_count = None):
+    """
+    Get the path to the JSON file(s) containing response data from the OpenAI
+    API.
+    """
+    save_path = Path(save_path)
+    os.makedirs(save_path, exist_ok=True)
+    
+    if chapter_index is not None:
+        save_path /= str(chapter_index)
+        if chapter_segment_index is not None:
+            save_path /= (f'{chapter_segment_index + 1}-of-'
+                          f'{chapter_segment_count}.json')
+    return save_path
+
+
+
+
+def save_openai_api_responses_for_script(script, prompt, encoding, max_chapter_segment_token_count, idx, api_key, model_id):
+    """
+    Call the OpenAI API for each chapter segment in a script and save the
+    responses to a list.
+    """
+    project_gutenberg_id = script['id']
+    chapter_count = len(script['chapters'])
+    logger.info(f'Starting to call OpenAI API and process responses for script '
+                f'{project_gutenberg_id} ({chapter_count} chapters).')
+
+    prompt_message_lists = []
+    response_list = []
+
+    for chapter in script['chapters']:
+        chapter_index = chapter['index']
+        chapter_segments = chapter['text']
+        chapter_segment_count = len(chapter_segments)
+
+        for chapter_segment_index, chapter_segment in enumerate(chapter_segments):
+            prompt_with_story = prompt.replace('{STORY}', chapter_segment)
+            prompt_message_lists.append([{
+                'role': 'user',
+                'content': prompt_with_story,
+                'api_key': api_key,
+                'model_id': model_id
+            }])
+
+    responses = execute_function_in_parallel(save_openai_api_response, prompt_message_lists)
+
+    for response in responses:
+        response_list.append(response)
+
+    logger.info(f'Finished processing responses for script {project_gutenberg_id}.')
+    return response_list
+
+
+def save_triples_for_scripts(input_data, idx, api_key, model_id):
+    """
+    Call the OpenAI API to generate knowledge graph nodes and edges, and store
+    the responses in a list.
+    """
+    # 1) load data
+    script = input_data
+
+    # 2) call OpenAI API
+    prompt = PROMPT_FILE_PATH.read_text()  # load prompt
+    max_chapter_segment_token_count = get_max_chapter_segment_token_count(prompt, model_id)
+    encoding = get_openai_model_encoding(model_id)
+    responses = save_openai_api_responses_for_script(
+        script, prompt, encoding, max_chapter_segment_token_count, idx, api_key, model_id
+    )
+
+    return responses

src/kg/utils.py

@@ -0,0 +1,57 @@
+import logging
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from pathlib import Path
+from typing import Callable
+
+LOGS_DIRECTORY_PATH = Path('../logs')
+# Number of processes to use when executing functions in parallel.
+MAX_PROCESS_COUNT = 20
+
+
+def set_up_logging(log_name):
+    """Set up a logger that logs to both the console and a file."""
+    log_path = LOGS_DIRECTORY_PATH / log_name
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+    logging.basicConfig(level=logging.INFO)
+    logger = logging.getLogger()
+    # Remove the default handler.
+    logger.handlers.clear()
+    stream_handler = logging.StreamHandler()
+    stream_handler.setLevel(logging.INFO)
+    stream_formatter = logging.Formatter(
+        '%(asctime)s - %(levelname)s - %(message)s', datefmt='%H:%M:%S')
+    stream_handler.setFormatter(stream_formatter)
+    logger.addHandler(stream_handler)
+    file_handler = logging.FileHandler(log_path)
+    file_handler.setLevel(logging.INFO)
+    file_formatter = logging.Formatter(
+        '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S')
+    file_handler.setFormatter(file_formatter)
+    logger.addHandler(file_handler)
+    return logger
+
+
+def strip_and_remove_empty_strings(strings):
+    """Strip a list of strings and remove empty strings."""
+    strings = [string.strip() for string in strings]
+    strings = [string for string in strings if string]
+    return strings
+
+
+def execute_function_in_parallel(function: Callable, *argument_lists,
+                                 logger=None):
+    """Execute a function in parallel using multiple processes."""
+    with ProcessPoolExecutor(max_workers=MAX_PROCESS_COUNT) as executor:
+        futures = [executor.submit(function, *arguments)
+                   for arguments in zip(*argument_lists)]
+        results = []
+        for future in as_completed(futures):
+            try:
+                result = future.result()
+                results.append(result)
+            except Exception as exception:
+                if logger:
+                    logger.exception('Exception')
+                raise exception
+        return results

src/summary/prompt.py

@@ -0,0 +1,33 @@
+import os
+from typing import List, Optional
+
+
+def build_summarizer_prompt(
+    prompt_template:str,
+    input_text_list:List[str],
+    chat_mode:Optional[str] = None) -> str:
+    
+    """_summary_
+    chat_mode(str) : 'hf-chat', 'kullm', 'None'
+
+    Returns:
+        _type_: _description_
+    """
+    
+    if os.path.isfile(prompt_template):
+        with open(prompt_template,'r') as f:
+            prompt_template = f.read() 
+    else:
+        pass
+    
+    # 예외처리 필요
+    assert isinstance(prompt_template, str)
+    
+    prompt = prompt_template.format(*input_text_list)
+    
+    if chat_mode == "hf-chat":
+        prompt = _get_hf_chat_template().format(prompt)
+    elif chat_mode == "kullm":
+        prompt = _get_kullm_template().format(prompt)
+    
+    return prompt

src/summary/summarizer.py

@@ -0,0 +1,65 @@
+
+import os
+import math
+from typing import Union, Optional
+import torch
+import logging
+
+#from vllm import LLM, SamplingParams
+#from vllm.lora.request import LoRARequest
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, set_seed, BitsAndBytesConfig
+import openai
+from openai.error import (APIError, RateLimitError, ServiceUnavailableError,
+                          Timeout, APIConnectionError, InvalidRequestError)
+from tenacity import (before_sleep_log, retry, retry_if_exception_type,
+                      stop_after_delay, wait_random_exponential, stop_after_attempt)
+
+
+logger = logging.getLogger(__name__)
+
+class Summarizer:
+    def __init__(self,
+                 inference_mode:str,
+                 model_id:str,
+                 api_key:str,
+                 dtype:str="bfloat16",
+                 seed=42,
+                 context_size:int=int(1024*26),
+                 gpu_memory_utilization:int=0.7,
+                 tensor_parallel_size=1
+                 ) -> None:
+        
+        self.inference_mode=inference_mode
+        self.model = None
+        self.tokenizer = None
+        self.seed = seed
+        openai.api_key = api_key
+        self.model = model_id
+    
+    def get_generation_config(
+        self,
+        repetition_penalty:float = 1.2,
+        do_sample:bool=True,
+        temperature:float = 0.1,
+        top_p:float = 0.9,
+        max_tokens:int = 1024
+        ):
+ 
+        return generation_config
+    
+    @retry(retry=retry_if_exception_type((APIError, Timeout, RateLimitError,
+                                        ServiceUnavailableError, APIConnectionError, InvalidRequestError)),
+        wait=wait_random_exponential(max=60), stop=stop_after_attempt(10),
+        before_sleep=before_sleep_log(logger, logging.WARNING))
+    def inference_with_gpt(self, prompt):
+        prompt_messages = [{"role": "user", "content": prompt}]
+        try:
+            response = openai.ChatCompletion.create(model = self.model, messages = prompt_messages, temperature = 0.1)
+            #finish_reason = response.choices[0].finish_reason
+            response = response.choices[0].message.content
+        except InvalidRequestError:
+            response = ''
+        
+        return response
+    

src/summary/utils.py

@@ -0,0 +1,118 @@
+#src.summary.utils.py
+import re
+from typing import List
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from typing import Callable
+
+from transformers import AutoTokenizer
+from tiktoken import Encoding, encoding_for_model
+
+SCENE_INDICATORS = ['씬/','씬','SS##','S#','s#','S','s','#\d+.','\d+.']
+
+def delete_special(pre_text, character_list):
+        for c in character_list:
+            pre_text = pre_text.replace(c, "")
+        return pre_text
+
+def preprocess_script(script:str) -> str:
+
+    lines = script.split("\n")
+    
+    new_text = ""
+    for line in lines:
+        line = delete_special(line, ["\n", "\t", "\xa0",'၀','ᝰ','ศ','ನ','tุ','\x00Ā\x00\x00\x00'])
+        cleaned = re.sub('[^가-힣a-zA-Z0-9\s,.!?/#]',' ', line).strip()
+        cleaned = delete_special(cleaned, ["  "]).strip()
+        cleaned = cleaned.replace("<|start|>", "").replace("<|end|>","")
+        if len(cleaned)>0:
+            new_text += f"{line}\n"
+    new_text = new_text.strip()
+    
+    return new_text
+
+
+def preprocess_scripts(scripts:List[str]) -> List[str]:
+    scripts = [preprocess_script(s) for s in scripts]
+
+    return scripts
+
+def break_down2scenes(text: str):
+    # Split the text using "s#" as the delimiter
+    scenes = re.split(r'(s#\d+)', text)
+    
+    # Remove empty elements from the split results
+    scenes = [scene for scene in scenes if scene.strip()]
+    
+    scenes_list = []
+    current_scene_number = None
+
+    for i in range(0, len(scenes), 2):  # Process the "s#" marker and corresponding text as pairs
+        scene_marker = scenes[i].strip()
+        scene_number = int(scene_marker.split('#')[1])  # Extract only the number
+        scene_text = scenes[i+1].strip() if i+1 < len(scenes) else ""
+
+        # Verify that the scene numbers are in the correct order
+        if current_scene_number is not None:
+            expected_scene_number = current_scene_number + 1
+            if scene_number != expected_scene_number:
+                raise ValueError(f"Unexpected scene number: {scene_number}, expected {expected_scene_number}")
+
+        # Save the scene number and text together
+        scenes_list.append({
+            'detected_scene_number': scene_number,
+            'text': f"{scene_marker}\n{scene_text}".strip()
+        })
+    return scenes_list
+
+
+def chunk_script_gpt(script:str,
+                    model_id:str,
+                    chunk_size:int=-1) -> List[str]:
+    if chunk_size == -1:
+        chunks = [script]
+        print("Single Inference Mode")
+        return chunks
+
+    encoding = encoding_for_model(model_id)
+    
+    scenes = break_down2scenes(script)
+    
+    len_scenes = len(scenes)
+
+    chunks = []
+    if len_scenes > 10:
+        print(f"Num of detected scenes : {len_scenes}")
+
+        chunk = ""
+        token_len_chunk = 0
+        for i, scene_data in enumerate(scenes):
+            scene = scene_data["text"].strip()
+            token_len_scene = len(encoding.encode_ordinary(scene))
+            if token_len_chunk + token_len_scene > chunk_size:
+                if token_len_chunk == 0:
+                    chunk += scene
+                    token_len_chunk += token_len_scene
+                else:
+                    chunks.append(chunk)
+                    chunk = scene
+                    token_len_chunk = token_len_scene
+            else:
+                chunk += scene
+                token_len_chunk += token_len_scene
+
+            if i == len_scenes-1:
+                chunks.append(chunk)
+    else:
+        print(f"No Detected Scenes ({len_scenes})")
+        tokenized_script = encoding.encode_ordinary(script)
+        token_len_script = len(tokenized_script)
+        for start in range(0,token_len_script,chunk_size):
+            if start + chunk_size >= token_len_script:
+                end = token_len_script+1
+            else:
+                end = start+chunk_size
+            
+            chunk = encoding.decode(tokenized_script[start:end])
+            chunks.append(chunk)
+    print(f"Num of chunks : {len(chunks)}")
+    return chunks

templates/atomic_fact.txt

@@ -0,0 +1,5 @@
+I will give you a summary from a chunk of movie script.
+Your task is to provide me with a list of atomic facts expressed in the given summary.
+Each atomic fact should be described in a name-only third-person format.
+Please separate each atomic fact with a `\n`.
+Summary: {}

templates/external_summary.txt

@@ -0,0 +1,8 @@
+This is a part of a script from a Movie. Read the following content carefully, then answer my question:
+{}
+The script has ended now.
+
+Please summarize the content:
+- Provide a detailed summary of the key characters' actions, emotions, and situations as reflected in the dialogue or context.
+- Clearly state the outcome of the events.
+- The summary should be between 2 to 5 sentences long.

templates/fact_score.txt

@@ -0,0 +1,8 @@
+Consider the given statement and the related scene.
+Indicate whether the statement is supported by the scene.
+Negation of a false statement should be considered supported.
+If the statement is true, output 1.
+If the statement is false, output the reason why it is false.
+Scene: {}
+Statement: {}
+Output:

templates/fact_score_kg.txt

@@ -0,0 +1,9 @@
+Consider the given statement, the related scene, and the relationship subgraph.
+Indicate whether the statement is supported by the scene and the relationship subgraph.
+Negation of a false statement should be considered supported.
+If the statement is true, output 1.
+If the statement is false, output the reason why it is false.
+Scene: {}
+Relationship Subgraph: {}
+Statement: {}
+Output:

templates/self_correction.txt

@@ -0,0 +1,8 @@
+Below is a part of the script from the titled movie.
+- Script: {}
+Based on the 'Statement to Revise' and 'Reason for Revision,' create a 'Revised Summary' of the 'Summary of the Script.'
+Keep the revised summary concise and similar in length to the original summary.
+Do not directly copy any part of the 'Script.'
+If the 'Summary of the Script' is accurate, generate the original summary as is.
+- Summary of the Script: {}
+- Revised Summary:

templates/story-prompt.txt

@@ -0,0 +1,43 @@
+Read part of a story, then identify named entities and generate knowledge graph edges.
+
+<Important Instructions>
+- Do not use or refer to the content, names, or entities from the example provided below in your response.
+- Only apply the process of identifying named entities and generating knowledge graph edges to the new story excerpt provided.
+
+<Example> (for reference only, do not use in your response):
+[Begin story excerpt]
+"Christmas won't be Christmas without any presents," grumbled Jo. "It's so dreadful to be poor!" sighed Meg, looking out the window at the snow-covered streets of Concord. "I don't think it's fair for some girls to have plenty of pretty things, and other girls nothing at all," added little Amy, with an injured sniff. "We've got Father and Mother, and each other," said Beth contentedly from her corner. The four young faces brightened at the cheerful words, but darkened again as Jo said sadly, "We haven't got Father, and shall not have him for a long time." She didn't say "perhaps never," but each silently added it, thinking of Father far away, where the fighting was.
+As young readers like to know 'how people look', we will take this moment to give them a little sketch of the four sisters. Margaret March, the eldest of the four, was sixteen, and very pretty, with large eyes, plenty of soft brown hair, a sweet mouth, and white hands. Fifteen-year-old Jo March was very tall, thin, and brown, and never seemed to know what to do with her long limbs. Elizabeth, or Beth, as everyone called her, was a rosy, smooth-haired, bright-eyed girl of thirteen, with a shy manner, a timid voice, and a peaceful expression which was seldom disturbed. Amy, the youngest, was a regular snow maiden, with blue eyes, and yellow hair curling on her shoulders.
+The clock struck six and, having swept up the hearth, Beth put a pair of slippers down to warm. Somehow the sight of the old shoes had a good effect upon the girls, for Mother was coming, and everyone brightened to welcome her. Jo sat up to hold the slippers nearer to the blaze. "They are quite worn out. Marmee must have a new pair." "I thought I'd get her some with my dollar," said Beth. "No, I shall!" cried Amy. "I'll tell you what we'll do," said Beth, "let's each get her something for Christmas, and not get anything for ourselves." "Let Marmee think we are getting things for ourselves, and then surprise her. We must go shopping tomorrow afternoon," said Jo, marching up and down.
+"Glad to find you so merry, my girls," said a cheery voice at the door, and the girls turned to welcome a tall, motherly lady. She was not elegantly dressed, but the girls thought the gray cloak and unfashionable bonnet covered the most splendid mother in the world. As they gathered about the table, Mrs. March said, with a particularly happy face, "I've got a treat for you after supper." A quick, bright smile went round like a streak of sunshine. Beth clapped her hands, and Jo tossed up her napkin, crying, "A letter! A letter! Three cheers for Father!" "Yes, a nice long letter. He is well, and he sends all sorts of loving wishes for Christmas, and an especial message to you girls," said Mrs. March, patting her pocket as if she had got a treasure there. "I think it was so splendid in Father to go as chaplain when he was too old to be drafted, and not strong enough for a soldier," said Meg warmly, proud of her father's work with the Union Army.
+[End story excerpt]
+
+Named entities (include all aliases and name variations):
+Jo / Jo March
+Meg / Margaret / Margaret March
+Amy
+Beth / Elizabeth
+March sisters
+Mrs. March / Marmee / Mother
+Father
+Concord
+Union Army
+
+Knowledge graph edges (select up to 20 most important, `subject(s); predicate; [object(s)]` format, named entities only, predicate(relation): two words or fewer):
+1. Jo, Meg, Amy, Beth; in; March sisters
+2. March sisters; daughters of; Mrs. March, Father
+3. Mrs. March; mother of; March sisters
+4. Father; father of; March sisters
+5. March sisters, Mrs. March; living in; Concord
+6. Father; fighting in war
+7. Father; chaplain in; Union Army
+8. Meg; sixteen years
+9. Jo; fifteen years
+10. Beth; thirteen years
+11. Beth; shy
+12. Amy; youngest among; March sisters
+
+<Request>
+[Begin story excerpt]
+{STORY}
+[End story excerpt]