Spaces:
Running
Running
| from utils.references import References | |
| from utils.file_operations import hash_name, make_archive, copy_templates | |
| from section_generator import section_generation_bg, keywords_generation | |
| import logging | |
| TOTAL_TOKENS = 0 | |
| TOTAL_PROMPTS_TOKENS = 0 | |
| TOTAL_COMPLETION_TOKENS = 0 | |
| 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 generate_backgrounds(title, description="", template="ICLR2022", model="gpt-4"): | |
| paper = {} | |
| paper_body = {} | |
| # Create a copy in the outputs folder. | |
| bibtex_path, destination_folder = copy_templates(template, title) | |
| logging.basicConfig(level=logging.INFO, filename=destination_folder + "/generation.log") | |
| # Generate keywords and references | |
| print("Initialize the paper information ...") | |
| input_dict = {"title": title, "description": description} | |
| keywords, usage = keywords_generation(input_dict, model="gpt-3.5-turbo") | |
| print(f"keywords: {keywords}") | |
| log_usage(usage, "keywords") | |
| ref = References(load_papers = "") | |
| ref.collect_papers(keywords, method="arxiv") | |
| 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() # to_prompts(top_papers) | |
| paper["body"] = paper_body | |
| paper["bibtex"] = bibtex_path | |
| for section in ["introduction", "related works", "backgrounds"]: | |
| try: | |
| # usage = pipeline(paper, section, destination_folder, model=model) | |
| usage = section_generation_bg(paper, section, destination_folder, 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 {destination_folder}.") | |
| # shutil.make_archive("output.zip", 'zip', save_to_path) | |
| input_dict = {"title": title, "description": description, "generator": "generate_backgrounds"} | |
| filename = hash_name(input_dict) + ".zip" | |
| return make_archive(destination_folder, filename) | |
| def fake_generator(title, description="", template="ICLR2022", model="gpt-4"): | |
| """ | |
| This function is used to test the whole pipeline without calling OpenAI API. | |
| """ | |
| input_dict = {"title": title, "description": description, "generator": "generate_backgrounds"} | |
| filename = hash_name(input_dict) + ".zip" | |
| return make_archive("sample-output.pdf", filename) | |