Bug fix. Put all generators in auto_backgrounds.py.

app.py

@@ -1,14 +1,14 @@
 import gradio as gr
 import os
 import openai
-from auto_backgrounds import generate_backgrounds, fake_generator
-from auto_draft import generate_draft
+from auto_backgrounds import generate_backgrounds, fake_generator, generate_draft
+from utils.file_operations import hash_name
 
 # todo:
 #   1. update README.md and introduction in app.py
 #   2. update QQ group and Organization cards
-#   3. update autodraft.py to generate a whole paper
 #   4. add auto_polishing function
+#   5. Use Completion to substitute some simple task (including: writing abstract, conclusion, generate keywords, generate figures...)
 
 openai_key = os.getenv("OPENAI_API_KEY")
 access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
@@ -44,16 +44,22 @@ def wrapped_generator(title, description, openai_key = None,
     #        if so, download from the cloud storage, return it
     #        if not, generate the result.
     if generator is None:
-        generator = generate_backgrounds
+        # todo: add a Dropdown to select which generator to use.
+        # generator = generate_backgrounds
+        # generator = generate_draft
+        generator = fake_generator
     if openai_key is not None:
         openai.api_key = openai_key
         openai.Model.list()
 
     if cache_mode:
-        from utils.storage import list_all_files, hash_name, download_file, upload_file
+        from utils.storage import list_all_files, download_file, upload_file
         # check if "title"+"description" have been generated before
-        file_name = hash_name(title, description) + ".zip"
+
+        input_dict = {"title": title, "description": description, "generator": "generate_draft"} #todo: modify here also
+        file_name = hash_name(input_dict) + ".zip"
         file_list = list_all_files()
+        # print(f"{file_name} will be generated. Check the file list {file_list}")
         if file_name in file_list:
             # download from the cloud storage, return it
             download_file(file_name)
@@ -61,12 +67,12 @@ def wrapped_generator(title, description, openai_key = None,
         else:
             # generate the result.
             # output = fake_generate_backgrounds(title, description, openai_key) # todo: use `generator` to control which function to use.
-            output = generate_backgrounds(title, description,  template, "gpt-4")
-            upload_file(file_name)
+            output = generator(title, description,  template, "gpt-4")
+            upload_file(output)
             return output
     else:
         # output = fake_generate_backgrounds(title, description, openai_key)
-        output = generate_backgrounds(title, description,  template, "gpt-4")
+        output = generator(title, description,  template, "gpt-4")
         return output
 
 
@@ -80,21 +86,22 @@ with gr.Blocks(theme=theme) as demo:
     gr.Markdown('''
     # Auto-Draft: 文献整理辅助工具
     
-    本Demo提供对[Auto-Draft](https://github.com/CCCBora/auto-draft)的auto_backgrounds功能的测试。通过输入一个领域的名称（比如Deep Reinforcement Learning)，即可自动对这个领域的相关文献进行归纳总结.    
+    本Demo提供对[Auto-Draft](https://github.com/CCCBora/auto-draft)的auto_draft功能的测试。通过输入想要生成的论文名称（比如Playing atari with deep reinforcement learning)，即可由AI辅助生成论文模板.    
     
-    ***2023-05-03 Update***: 在这个版本中为大家提供了输入OpenAI API Key的地址, 如果有GPT-4的API KEY的话可以在这里体验! 
+    ***2023-05-03 Update***: 在公开版本中为大家提供了输入OpenAI API Key的地址, 如果有GPT-4的API KEY的话可以在这里体验! 
     我也会在近期提供一定的免费体验在这个Huggingface Organization里： [AUTO-ACADEMIC](https://huggingface.co/organizations/auto-academic/share/HPjgazDSlkwLNCWKiAiZoYtXaJIatkWDYM).
     如果有更多想法和建议欢迎加入QQ群里交流, 如果我在Space里更新了Key我会第一时间通知大家. 群号: ***249738228***.  
     
     ## 用法
     
-    输入一个领域的名称（比如Deep Reinforcement Learning), 点击Submit, 等待大概十分钟, 下载.zip格式的输出，在Overleaf上编译浏览.  
+    输入想要生成的论文名称（比如Playing Atari with Deep Reinforcement Learning), 点击Submit, 等待大概十分钟, 下载.zip格式的输出，在Overleaf上编译浏览.  
     ''')
     with gr.Row():
         with gr.Column(scale=2):
             key =  gr.Textbox(value=openai_key, lines=1, max_lines=1, label="OpenAI Key", visible=not IS_OPENAI_API_KEY_AVAILABLE)
-            title = gr.Textbox(value="Deep Reinforcement Learning", lines=1, max_lines=1, label="Title")
-            description = gr.Textbox(lines=5, label="Description (Optional)")
+            # generator = gr.Dropdown(choices=["学术论文", "文献总结"], value="文献总结", label="Selection", info="目前支持生成'学术论文'和'文献总结'.", interactive=True)
+            title = gr.Textbox(value="Playing Atari with Deep Reinforcement Learning", lines=1, max_lines=1, label="Title", info="论文标题")
+            description = gr.Textbox(lines=5, label="Description (Optional)", visible=False)
 
             with gr.Row():
                 clear_button = gr.Button("Clear")
@@ -104,8 +111,8 @@ with gr.Blocks(theme=theme) as demo:
             availability_mapping = {True: "AVAILABLE", False: "NOT AVAILABLE"}
             gr.Markdown(f'''## Huggingface Space Status  
              当`OpenAI API`显示AVAILABLE的时候这个Space可以直接使用.    
-             当`OpenAI API`显示NOT AVAILABLE的时候这个Space可以通过在左侧输入OPENAI KEY来使用. 需要有GPT-4的API权限, 不然会报错. 
-             当`Cache`显示AVAILABLE的时候, 所有的输入和输出会被备份到我的云储存中. 显示NOT AVAILABLE的时候可以正常使用.
+             当`OpenAI API`显示NOT AVAILABLE的时候这个Space可以通过在左侧输入OPENAI KEY来使用. 需要有GPT-4的API权限. 
+             当`Cache`显示AVAILABLE的时候, 所有的输入和输出会被备份到我的云储存中. 显示NOT AVAILABLE的时候不影响实际使用.
             `OpenAI API`: <span style="{style_mapping[IS_OPENAI_API_KEY_AVAILABLE]}">{availability_mapping[IS_OPENAI_API_KEY_AVAILABLE]}</span>.  `Cache`: <span style="{style_mapping[IS_CACHE_AVAILABLE]}">{availability_mapping[IS_CACHE_AVAILABLE]}</span>.''')
             file_output = gr.File(label="Output")
 

auto_backgrounds.py

@@ -1,12 +1,13 @@
 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
+from section_generator import section_generation_bg, keywords_generation, figures_generation, section_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
@@ -26,7 +27,7 @@ def log_usage(usage, generating_target, print_out=True):
         print(message)
     logging.info(message)
 
-def generate_backgrounds(title, description="", template="ICLR2022", model="gpt-4"):
+def _generation_setup(title, description="", template="ICLR2022", model="gpt-4"):
     paper = {}
     paper_body = {}
 
@@ -41,9 +42,9 @@ def generate_backgrounds(title, description="", template="ICLR2022", model="gpt-
     print(f"keywords: {keywords}")
     log_usage(usage, "keywords")
 
-    ref = References(load_papers = "")
+    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
+    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}.")
 
@@ -52,6 +53,12 @@ def generate_backgrounds(title, description="", template="ICLR2022", model="gpt-
     paper["references"] = ref.to_prompts()
     paper["body"] = paper_body
     paper["bibtex"] = bibtex_path
+    return paper, destination_folder, all_paper_ids
+
+
+
+def generate_backgrounds(title, description="", template="ICLR2022", model="gpt-4"):
+    paper, destination_folder, _ = _generation_setup(title, description, template, model)
 
     for section in ["introduction", "related works", "backgrounds"]:
         try:
@@ -70,6 +77,31 @@ 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"}
+    input_dict = {"title": title, "description": description, "generator": "generate_draft"}
     filename = hash_name(input_dict) + ".zip"
     return make_archive("sample-output.pdf", filename)
+
+
+def generate_draft(title, description="", template="ICLR2022", model="gpt-4"):
+    paper, destination_folder, _ = _generation_setup(title, description, template, model)
+
+    print("Generating figures ...")
+    usage = figures_generation(paper, destination_folder, model="gpt-3.5-turbo")
+    # todo: use `figures_generation` function to complete remainings
+    # 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"]:
+    for section in ["introduction", "related works", "backgrounds", "experiments", "conclusion", "abstract"]:
+        try:
+            usage = section_generation(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}")
+
+    input_dict = {"title": title, "description": description, "generator": "generate_draft"}
+    filename = hash_name(input_dict) + ".zip"
+    return make_archive(destination_folder, filename)

auto_draft.py

@@ -1,145 +1,145 @@
-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)
+# 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)

section_generator.py

@@ -1,5 +1,6 @@
 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
 
@@ -43,6 +44,38 @@ def section_generation_bg(paper, section, save_to_path, model):
     return usage
 
 
+def section_generation(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 = 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 keywords_generation(input_dict,  model):
     title = input_dict.get("title")
     description = input_dict.get("description", "")
@@ -54,5 +87,9 @@ def keywords_generation(input_dict,  model):
     else:
         raise ValueError("`input_dict` must include the key 'title'.")
 
-def figures_generation():
-    pass
+def figures_generation(paper, save_to_path, model):
+    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

utils/storage.py

@@ -1,45 +1,52 @@
 # This script `storage.py` is used to handle the cloud storage.
 #   `upload_file`:
+#       Function to upload a local file to the specified S3 bucket.
+#       If the target_name is not specified, it will use the file_name as the object key.
 #   `list_all_files`:
+#       Function to list all the files in the specified S3 bucket.
 #   `download_file`:
+#       Function to download a file from the specified S3 bucket to the local machine using the specified file_name.
 
 import os
 import boto3
 
-access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
-secret_access_key = os.getenv('AWS_SECRET_ACCESS_KEY')
-bucket_name = "hf-storage"
+BUCKET_NAME = "hf-storage"
 
-if (access_key_id is not None) and (secret_access_key is not None):
+def get_client():
+    access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
+    secret_access_key = os.getenv('AWS_SECRET_ACCESS_KEY')
     session = boto3.Session(
         aws_access_key_id=access_key_id,
         aws_secret_access_key=secret_access_key,
     )
-
     s3 = session.resource('s3')
-    bucket = s3.Bucket(bucket_name)
-
-
-    def upload_file(file_name, target_name=None):
-        if target_name is None:
-            target_name = file_name
-        try:
-            s3.meta.client.upload_file(Filename=file_name, Bucket=bucket_name, Key=target_name)
-            print(f"The file {file_name} has been uploaded!")
-        except:
-            print("Uploading failed!")
-
-    def list_all_files():
-        return [obj.key for obj in bucket.objects.all()]
-
-    def download_file(file_name):
-        ''' Download `file_name` from the bucket.
-        Bucket (str) – The name of the bucket to download from.
-        Key (str) – The name of the key to download from.
-        Filename (str) – The path to the file to download to.
-        '''
-        try:
-            s3.meta.client.download_file(Bucket=bucket_name, Key=file_name, Filename=file_name)
-            print(f"The file {file_name} has been downloaded!")
-        except:
-            print("Uploading failed!")
+    bucket = s3.Bucket(BUCKET_NAME)
+    return s3, bucket
+
+def upload_file(file_name, target_name=None):
+    s3, _ = get_client()
+
+    if target_name is None:
+        target_name = file_name
+    s3.meta.client.upload_file(Filename=file_name, Bucket=BUCKET_NAME, Key=target_name)
+    print(f"The file {file_name} has been uploaded!")
+
+
+def list_all_files():
+    _, bucket = get_client()
+    return [obj.key for obj in bucket.objects.all()]
+
+
+def download_file(file_name):
+    ''' Download `file_name` from the bucket.
+    Bucket (str) – The name of the bucket to download from.
+    Key (str) – The name of the key to download from.
+    Filename (str) – The path to the file to download to.
+    '''
+    s3, _ = get_client()
+    s3.meta.client.download_file(Bucket=BUCKET_NAME, Key=file_name, Filename=file_name)
+    print(f"The file {file_name} has been downloaded!")
+
+if __name__ == "__main__":
+    file = "sample-output.pdf"
+    upload_file(file)