File size: 4,338 Bytes
a0d1776
 
ae495a3
a0d1776
 
2dc9347
c9efba3
 
a0d1776
 
 
 
 
 
 
2dc9347
a0d1776
 
 
 
 
 
 
 
 
 
 
 
 
c9efba3
a0d1776
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2dc9347
ae495a3
 
 
 
 
 
 
 
 
2dc9347
 
ae495a3
 
2dc9347
 
05783f8
ae495a3
 
c9efba3
2dc9347
 
 
 
 
 
 
 
 
 
 
 
a0d1776
c9efba3
2dc9347
c9efba3
c160ff7
c9efba3
 
 
72c76c9
c9efba3
 
2dc9347
 
 
 
 
 
 
a0d1776
ae495a3
2dc9347
ae495a3
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
from utils.prompts import generate_paper_prompts, generate_keywords_prompts, generate_experiments_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):
    """
    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)
    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") 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