section_generator.py

# from utils.prompts import generate_paper_prompts, generate_keywords_prompts, generate_experiments_prompts, generate_bg_summary_prompts
from utils.prompts import generate_paper_prompts, generate_bg_summary_prompts
# from utils.gpt_interaction import get_responses #, extract_responses, extract_keywords, extract_json
from utils.figures import generate_random_figures
import time
import os
from utils.prompts import KEYWORDS_SYSTEM, SECTION_GENERATION_SYSTEM
from utils.gpt_interaction import get_gpt_responses
import json

#  three GPT-based content generator:
#       1. section_generation: used to generate main content of the paper
#       2. keywords_generation: used to generate a json output {key1: output1, key2: output2} for multiple purpose.
#       3. figure_generation: used to generate sample figures.
#  all generator should return the token usage.

MAX_ATTEMPTS = 6


def section_generation_bg(paper, section, save_to_path, model):
    """
    todo: this part should be revised
    The main pipeline of generating a section.
        1. Generate prompts.
        2. Get responses from AI assistant.
        3. Extract the section text.
        4. Save the text to .tex file.
    :return usage
    """
    print(f"Generating {section}...")
    prompts = generate_bg_summary_prompts(paper, section)
    # gpt_response, usage = get_responses(prompts, model)
    gpt_response, usage = get_gpt_responses(prompts, model)
    output = gpt_response  # extract_responses(gpt_response)
    paper["body"][section] = output
    tex_file = os.path.join(save_to_path, f"{section}.tex")
    # tex_file = save_to_path + f"/{section}.tex"
    if section == "abstract":
        with open(tex_file, "w") as f:
            f.write(r"\begin{abstract}")
        with open(tex_file, "a") as f:
            f.write(output)
        with open(tex_file, "a") as f:
            f.write(r"\end{abstract}")
    else:
        with open(tex_file, "w") as f:
            f.write(f"\section{{{section.upper()}}}\n")
        with open(tex_file, "a") as f:
            f.write(output)
    time.sleep(5)
    print(f"{section} has been generated. Saved to {tex_file}.")
    return usage


def section_generation(paper, section, save_to_path, model, research_field="machine learning"):
    """
    The main pipeline of generating a section.
        1. Generate prompts.
        2. Get responses from AI assistant.
        3. Extract the section text.
        4. Save the text to .tex file.
    :return usage
    """
    prompts = generate_paper_prompts(paper, section)
    output, usage = get_gpt_responses(SECTION_GENERATION_SYSTEM.format(research_field=research_field), prompts,
                                      model=model, temperature=0.4)
    paper["body"][section] = output
    tex_file = os.path.join(save_to_path, f"{section}.tex")
    with open(tex_file, "w", encoding="utf-8") as f:
        f.write(output)
    time.sleep(5)
    return usage


def keywords_generation(input_dict, default_keywords=None):
    """
    Input:
        input_dict: a dictionary containing the title of a paper.
        default_keywords: if anything went wrong, return this keywords.

    Output:
        a dictionary including all keywords and their importance score.

    Input example: {"title": "The title of a Machine Learning Paper"}
    Output Example: {"machine learning": 5, "reinforcement learning": 2}
    """
    title = input_dict.get("title")
    attempts_count = 0
    while (attempts_count < MAX_ATTEMPTS) and (title is not None):
        try:
            keywords, usage = get_gpt_responses(KEYWORDS_SYSTEM.format(min_refs_num=1, max_refs_num=10), title,
                                                model="gpt-3.5-turbo", temperature=0.4)
            print(keywords)
            output = json.loads(keywords)
            return output.keys(), usage
        except json.decoder.JSONDecodeError:
            attempts_count += 1
            time.sleep(10)
    # Default references
    print("Error: Keywords generation has failed. Return the default keywords.")
    if default_keywords is None or isinstance(default_keywords, dict):
        return {"machine learning": 10}
    else:
        return default_keywords

# def figures_generation(paper, save_to_path, model):
#     # todo: this function is not complete.
#     prompts = generate_experiments_prompts(paper)
#     gpt_response, usage = get_responses(prompts, model)
#     list_of_methods = list(extract_json(gpt_response))
#     generate_random_figures(list_of_methods, os.path.join(save_to_path, "comparison.png"))
#     return usage