auto-draft / auto_draft.py
sc_ma
Bug fix. Put all generators in auto_backgrounds.py.
ae495a3
raw
history blame
5.87 kB
# from utils.references import References
# from utils.prompts import generate_paper_prompts, generate_keywords_prompts, generate_experiments_prompts
# from utils.gpt_interaction import get_responses, extract_responses, extract_keywords, extract_json
# from utils.tex_processing import replace_title
# from utils.figures import generate_random_figures
# import datetime
# import shutil
# import time
# import logging
# import os
#
# TOTAL_TOKENS = 0
# TOTAL_PROMPTS_TOKENS = 0
# TOTAL_COMPLETION_TOKENS = 0
#
# def make_archive(source, destination):
# base = os.path.basename(destination)
# name = base.split('.')[0]
# format = base.split('.')[1]
# archive_from = os.path.dirname(source)
# archive_to = os.path.basename(source.strip(os.sep))
# shutil.make_archive(name, format, archive_from, archive_to)
# shutil.move('%s.%s'%(name,format), destination)
# return destination
#
#
# def log_usage(usage, generating_target, print_out=True):
# global TOTAL_TOKENS
# global TOTAL_PROMPTS_TOKENS
# global TOTAL_COMPLETION_TOKENS
#
# prompts_tokens = usage['prompt_tokens']
# completion_tokens = usage['completion_tokens']
# total_tokens = usage['total_tokens']
#
# TOTAL_TOKENS += total_tokens
# TOTAL_PROMPTS_TOKENS += prompts_tokens
# TOTAL_COMPLETION_TOKENS += completion_tokens
#
# message = f"For generating {generating_target}, {total_tokens} tokens have been used ({prompts_tokens} for prompts; {completion_tokens} for completion). " \
# f"{TOTAL_TOKENS} tokens have been used in total."
# if print_out:
# print(message)
# logging.info(message)
#
# def pipeline(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_paper_prompts(paper, section)
# gpt_response, usage = get_responses(prompts, model)
# output = extract_responses(gpt_response)
# paper["body"][section] = output
# 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}}}\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 generate_draft(title, description="", template="ICLR2022", model="gpt-4"):
# """
# The main pipeline of generating a paper.
# 1. Copy everything to the output folder.
# 2. Create references.
# 3. Generate each section using `pipeline`.
# 4. Post-processing: check common errors, fill the title, ...
# """
# paper = {}
# paper_body = {}
#
# # Create a copy in the outputs folder.
# # todo: use copy_templates function instead.
# now = datetime.datetime.now()
# target_name = now.strftime("outputs_%Y%m%d_%H%M%S")
# source_folder = f"latex_templates/{template}"
# destination_folder = f"outputs/{target_name}"
# shutil.copytree(source_folder, destination_folder)
#
# bibtex_path = destination_folder + "/ref.bib"
# save_to_path = destination_folder +"/"
# replace_title(save_to_path, title)
# logging.basicConfig( level=logging.INFO, filename=save_to_path+"generation.log")
#
# # Generate keywords and references
# print("Initialize the paper information ...")
# prompts = generate_keywords_prompts(title, description)
# gpt_response, usage = get_responses(prompts, model)
# keywords = extract_keywords(gpt_response)
# log_usage(usage, "keywords")
# ref = References(load_papers = "") #todo: allow users to upload bibfile.
# ref.collect_papers(keywords, method="arxiv") #todo: add more methods to find related papers
# all_paper_ids = ref.to_bibtex(bibtex_path) #todo: this will used to check if all citations are in this list
#
# print(f"The paper information has been initialized. References are saved to {bibtex_path}.")
#
# paper["title"] = title
# paper["description"] = description
# paper["references"] = ref.to_prompts() #todo: see if this prompts can be compressed.
# paper["body"] = paper_body
# paper["bibtex"] = bibtex_path
#
# print("Generating figures ...")
# prompts = generate_experiments_prompts(paper)
# gpt_response, usage = get_responses(prompts, model)
# list_of_methods = list(extract_json(gpt_response))
# log_usage(usage, "figures")
# generate_random_figures(list_of_methods, save_to_path + "comparison.png")
#
# for section in ["introduction", "related works", "backgrounds", "methodology", "experiments", "conclusion", "abstract"]:
# try:
# usage = pipeline(paper, section, save_to_path, model=model)
# log_usage(usage, section)
# except Exception as e:
# print(f"Failed to generate {section} due to the error: {e}")
# print(f"The paper {title} has been generated. Saved to {save_to_path}.")
# return make_archive(destination_folder, "output.zip")
#
# if __name__ == "__main__":
# # title = "Training Adversarial Generative Neural Network with Adaptive Dropout Rate"
# title = "Playing Atari Game with Deep Reinforcement Learning"
# description = ""
# template = "ICLR2022"
# model = "gpt-4"
# # model = "gpt-3.5-turbo"
#
# generate_draft(title, description, template, model)