plot_extract/summariser.py

import os
import json
import json_repair
import sys
import argparse
from preprocessor import preprocess_remove_line_break

project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)

from chatbot_api import llm_factory


parser = argparse.ArgumentParser()
parser.add_argument("--source-file-path", "-s", type=str, default=r"novels/海底两万里/Twenty Thousand Leagues Under the Seas.txt")
args = parser.parse_args()

source_file_path = args.source_file_path
if os.path.isabs(source_file_path):
    source_file_path = os.path.abspath(source_file_path)
else:
    source_file_path = os.path.join(project_root, source_file_path)
source_dir, source_file = os.path.split(source_file_path)
source_filename, _ = os.path.splitext(source_file)
output_filename = source_filename.replace(" ", "_").lower()
summarised_filename = os.path.join(source_dir, f"{output_filename}_summarised.json")

with open(source_file_path, 'r', encoding='utf-8') as f:
    full_novel_text = preprocess_remove_line_break(f.read(), " ")

chunk_size = 50000
model_type = "qwen"
model = "Qwen/Qwen3-0.6B"

llm = llm_factory(model_type=model_type, model_name=model, max_new_tokens=4096, temperature=0.1)

summariser_sys_prompt = """
You are a story summariser. You are tasked to summarise a long novel provided in chunks into a brief summary. The summary should include the main plot, characters, and all the major events.
Your output should be a json data object with the following structure:
{
    "reasoning": "The reasoning of the summary",
    "summary": "The summary of the novel",
    "tableofcontents_gold": [
        "The table of contents of the novel",
        ...
    ]
    "tableofcontents_custom": [
        "The table of contents of the novel",
        ...
    ]
}
And in the input you will also be given a json data object of the same structure, produced by the previous step.

The "reasoning" is your analysis and outputreasoning of the summary.
"summary" is the plain language summary of the novel. You may always rewrite, modify this and extend the last summary to make it more complete, and to include the content provided in the new chunks. The overall length of summary should not exceed 800 words. You may rewrite previous parts in more concise languages when necessary.
If there is an official table of contents provided in the provided novel, you extract it in the form of an array of strings in the "tableofcontents_gold" field, one chapter per element. "tableofcontents_gold" should be complete once filled.If "tableofcontents_gold" is filled already, do not modify it. If there is no official table of contents, leave "tableofcontents_gold" an empty array.
"tableofcontents_custom" is your custom table of contents extracted from the main body of the novel, regardless if there is an official table of contents provided in the provided novel or not. You may name and divide chapters according to the content of the novel. You may append and modify it based on the content provided in the new chunks. Every chapter need to contain a numerical index.
You should always provide all fields in the output json data object. And these output should always be produced for the whole novel, including what was provided in the previous chunks.
"""

summariser_user_prompt = """
Here is the json data object of the same structure, produced by the previous step:

{previous_summary}

Here is the novel:

{novel_text}
"""

head = 0
tail = min(head + chunk_size, len(full_novel_text))

previous_summary = "None"
gold_toc_extracted = False
while tail < len(full_novel_text):
    chunk = full_novel_text[head:tail]
    
    json_prefilled = "{\n   \"reasoning\":"
    messages = [
        {
            "role": "system",
            "content": summariser_sys_prompt
        },
        {
            "role": "user",
            "content": summariser_user_prompt.format(previous_summary=json.dumps(previous_summary, indent=4), novel_text=chunk)
        },
        {
            "role": "assistant",
            "content": json_prefilled
        }
    ]
    response = llm.generate(
        messages,
    )
    response_json_str = json_prefilled + response
    response_json = json_repair.loads(response_json_str)
    print(json.dumps(response_json, indent=4))
    response_json_no_reasoning = {
        "summary": response_json["summary"],
        "tableofcontents_gold": response_json["tableofcontents_gold"],
        "tableofcontents_custom": response_json["tableofcontents_custom"]
    }
    previous_summary = response_json_no_reasoning
    
    
    if gold_toc_extracted:
        response_json["tableofcontents_gold"] = previous_summary["tableofcontents_gold"]
    if "tableofcontents_gold" in response_json and response_json["tableofcontents_gold"]:
        gold_toc_extracted = True
    
    with open(summarised_filename, "w", encoding='utf-8') as f:
        json.dump(response_json_no_reasoning, f, indent=4)
    
    head = tail - chunk_size // 10
    tail = min(head + chunk_size, len(full_novel_text))