fix bug, reduce memory footprint

app.py

@@ -21,8 +21,9 @@ def sanitize_jinja2(jinja_lines):
 def get_existing_templates():
     return [None] + os.listdir("./templates")
 
-if len(os.listdir("./tmp")) > 20:
-    shutil.rmtree('./tmp')
+# if os.path.exists("./tmp"):
+#     if len(os.listdir("./tmp")) > 20:
+#         shutil.rmtree('./tmp')
 
 # Initialization
 if 'tokenizer_json' not in st.session_state:
@@ -83,12 +84,14 @@ if gen_button:
         st.session_state['tokenizer'] = AutoTokenizer.from_pretrained(hf_model_repo_name)
 
         st.session_state['repo_normalized_name'] = hf_model_repo_name.replace("/", "_")
-        st.session_state['tokenizer'].save_pretrained(f"./tmp/{st.session_state['uuid']}_{hf_model_repo_name}")
         st.session_state['tokenizer_json'] = f"./tmp/{st.session_state['uuid']}_{hf_model_repo_name}"
+        # st.session_state['tokenizer'].save_pretrained(st.session_state['tokenizer_json'])
     
 if st.session_state['tokenizer_json'] is not None:
+    st.session_state['tokenizer'].save_pretrained(st.session_state['tokenizer_json'])
     with open(f"{st.session_state['tokenizer_json']}/tokenizer_config.json", "rb") as f:
         tokenizer_json = json.load(f)
+    shutil.rmtree(st.session_state['tokenizer_json'])
 
     json_spec, col2 = st.columns(spec=[0.3, 0.7])
 
@@ -146,6 +149,7 @@ if st.session_state['tokenizer_json'] is not None:
                 with open(f"./tmp/{st.session_state['uuid']}/tmp_chat_template.json", "r") as f:
                     jinja_lines = f.readlines()
                     st.session_state['tokenizer'].chat_template = sanitize_jinja2(jinja_lines)
+                os.remove(f"./tmp/{st.session_state['uuid']}/tmp_chat_template.json")
                 generated_prompt_wo_add_generation_prompt = st.session_state['tokenizer'].apply_chat_template(chat, tokenize=False, add_generation_prompt= False)
                 generated_prompt_w_add_generation_prompt = st.session_state['tokenizer'].apply_chat_template(chat, tokenize=False, add_generation_prompt= True)
 
@@ -166,10 +170,11 @@ if st.session_state['tokenizer_json'] is not None:
                 to_private_checkbox = st.checkbox("To Private Repo", key="to_private_checkbox")
                 create_pr_checkbox = st.checkbox("Create PR (For Contribution 🤗)", key="create_pr_checkbox")
                 push_to_hub_button = st.button("Push to Hub", key="push_to_hub_button", use_container_width=True)
-                st.session_state['tokenizer'].save_pretrained(f"./tmp/{st.session_state['uuid']}_{hf_model_repo_name}")
-                with open(f"./tmp/{st.session_state['uuid']}_{hf_model_repo_name}/tokenizer_config.json", "r") as f:
+                st.session_state['tokenizer'].save_pretrained(st.session_state['tokenizer_json'])
+                with open(f"{st.session_state['tokenizer_json']}/tokenizer_config.json", "r") as f:
 
                     tokenizer_config_content = json.loads(f.read())
+                shutil.rmtree(st.session_state['tokenizer_json'])
 
                 st.download_button(
                     label="Download tokenizer_config.json",
@@ -195,3 +200,4 @@ if st.session_state['tokenizer_json'] is not None:
                         except Exception as e:
                             st.write(f"Repo id: {st.session_state['repo_id']}")
                             st.write(str(e))
+                    os.remove(f"./tmp/{st.session_state['uuid']}/tmp_chat_template.json")