src/kg/save_triples.py

# 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