Update app.py

app.py

@@ -17,7 +17,7 @@ import os
 import re
 st.set_page_config(layout="wide")
 
-@st.cache
+@st.cache_data
 def get_args():
     # -- 1. Setup arguments
     parser = argparse.ArgumentParser()
@@ -34,17 +34,15 @@ def get_args():
 
     # -- 3. Setup constants
     args = parser.parse_args()
-    return args
+    return args, logger
 
-@st.cache
-def get_podcast_data(transcription_path):
-    together.api_key = os.environ["TOGETHER_API_KEY"]
-    together.Models.start(MODEL)
-    podcast_url_video_df = pd.read_csv(PODCAST_URL_VIDEO_PATH, sep=';')
+@st.cache_data
+def get_podcast_data(path):
+    podcast_url_video_df = pd.read_csv(path, sep=';')
     return podcast_url_video_df
 
-@st.cache
-def setup_basics_comp(emb_model, model, default_system_prompt_link, logger, podcast_url_video_df, img_size=100):
+@st.cache_resource(experimental_allow_widgets=True)
+def get_basics_comp(emb_model, model, default_system_prompt_link, _logger, podcast_url_video_df, img_size=100):
     r    = requests.get("https://raw.githubusercontent.com/AlbertoUAH/Castena/main/media/castena-animated-icon.gif", stream=True)
     icon = Image.open(r.raw)
     icon = icon.resize((img_size, img_size))
@@ -56,17 +54,22 @@ def setup_basics_comp(emb_model, model, default_system_prompt_link, logger, podc
     video_option_joined = '_'.join(video_option.replace(': Entrevista a ', ' ').lower().split(' ')).replace("\'", "")
     video_option_joined_path = "{}_transcription.txt".format(video_option_joined)
     youtube_video_url   = list(podcast_url_video_df[podcast_url_video_df['podcast_name'].str.contains(video_option_joined)]['youtube_video_url'])[0].replace("\'", "")
+    st.title("[Podcast: {}]({})".format(video_option.replace("'", "").title(), youtube_video_url))
 
     # -- 4. Setup request for system prompt
     f = urllib.request.urlopen(default_system_prompt_link)
-    DEFAULT_SYSTEM_PROMPT = str(f.read(), 'UTF-8')
+    default_system_prompt = str(f.read(), 'UTF-8')
 
     # -- 5. Setup app
-    translator, nlp, retriever = utils.setup_app(video_option_joined_path, emb_model, model, logger)
-    return translator, nlp, retriever, video_option, video_option_joined_path
+    translator, nlp, retriever = utils.setup_app(video_option_joined_path, emb_model, model, _logger)
+
+    # -- 6. Setup model
+    together.api_key = os.environ["TOGETHER_API_KEY"]
+    together.Models.start(model)
+    return together, translator, nlp, retriever, video_option, video_option_joined_path, default_system_prompt, youtube_video_url
     
 def main():
-    args = get_args()
+    args, logger = get_args()
     B_INST, E_INST = "[INST]", "[/INST]"
     B_SYS, E_SYS   = "<<SYS>>\n", "\n<</SYS>>\n\n"
 
@@ -76,12 +79,14 @@ def main():
     TRANSCRIPTION_PATH         = '{}_transcription.txt'.format(TRANSCRIPTION)
     MODEL                      = args.MODEL
     EMB_MODEL                  = args.EMB_MODEL
+    WIDTH                      = 50
+    SIDE                       = (100 - WIDTH) / 2
 
-    podcast_url_video_df = get_podcast_data(TRANSCRIPTION_PATH)
+    podcast_url_video_df = get_podcast_data(PODCAST_URL_VIDEO_PATH)
 
-    translator, nlp, retriever, video_option, video_option_joined_path = setup_basics_comp(EMB_MODEL, MODEL, 
-                                                                                           DEFAULT_SYSTEM_PROMPT_LINK, logger, 
-                                                                                           podcast_url_video_df, img_size=100)
+    together, translator, nlp, retriever, video_option, video_option_joined_path, default_system_prompt, youtube_video_url = get_basics_comp(EMB_MODEL, MODEL, 
+                                                                                                                                             DEFAULT_SYSTEM_PROMPT_LINK, logger, 
+                                                                                                                                             podcast_url_video_df, img_size=100)
 
 
     # -- 6. Setup prompt template + llm chain
@@ -90,7 +95,7 @@ def main():
 PREGUNTA: {question}
 
 RESPUESTA: """
-    prompt_template = utils.get_prompt(instruction, DEFAULT_SYSTEM_PROMPT, B_SYS, E_SYS, B_INST, E_INST, logger)
+    prompt_template = utils.get_prompt(instruction, default_system_prompt, B_SYS, E_SYS, B_INST, E_INST, logger)
 
     llama_prompt = PromptTemplate(
         template=prompt_template, input_variables=["context", "question"]
@@ -100,12 +105,7 @@ RESPUESTA: """
     qa_chain = utils.create_llm_chain(MODEL, retriever, chain_type_kwargs, logger, video_option_joined_path)
 
     # ---------------------------------------------------------------------
-    # -- 7. Setup Streamlit app
-    st.title("[Podcast: {}]({})".format(video_option.replace("'", "").title(), youtube_video_url))
-
-    width = 50
-    side = (100 - width) / 2
-    _, container, _ = st.columns([side, width, side])
+    _, container, _ = st.columns([SIDE, WIDTH, SIDE])
     with container:
         st_player(utils.typewrite(youtube_video_url))
 
@@ -131,11 +131,9 @@ RESPUESTA: """
                     end_time_seconds_list.append(end_time_seconds)
 
             if start_time_str_list:
-                width = 40
-                side = (100 - width) / 2
                 for start_time_seconds, start_time_str, end_time_seconds in zip(start_time_seconds_list, start_time_str_list, end_time_seconds_list):
                     st.markdown("__Fragmento: " + start_time_str + "__")
-                    _, container, _ = st.columns([side, width, side])
+                    _, container, _ = st.columns([SIDE, WIDTH, SIDE])
                     with container:
                         st_player(youtube_video_url.replace("?enablejsapi=1", "") + f'?start={start_time_seconds}&end={end_time_seconds}')