Update app.py

app.py

@@ -15,11 +15,10 @@ import spacy
 import time
 import os
 import re
+st.set_page_config(layout="wide")
 
 @st.cache
 def get_args():
-    st.set_page_config(layout="wide")
-
     # -- 1. Setup arguments
     parser = argparse.ArgumentParser()
     parser.add_argument('--DEFAULT_SYSTEM_PROMPT_LINK', type=str, default="https://raw.githubusercontent.com/AlbertoUAH/Castena/main/prompts/default_system_prompt.txt", help='Valor para DEFAULT_SYSTEM_PROMPT_LINK')
@@ -43,25 +42,12 @@ def get_podcast_data(transcription_path):
     together.Models.start(MODEL)
     podcast_url_video_df = pd.read_csv(PODCAST_URL_VIDEO_PATH, sep=';')
     return podcast_url_video_df
-    
-def main():
-    args = get_args()
-    B_INST, E_INST = "[INST]", "[/INST]"
-    B_SYS, E_SYS   = "<<SYS>>\n", "\n<</SYS>>\n\n"
-
-    # -- 4. Get parameters
-    PODCAST_URL_VIDEO_PATH     = args.PODCAST_URL_VIDEO_PATH
-    DEFAULT_SYSTEM_PROMPT_LINK = args.DEFAULT_SYSTEM_PROMPT_LINK
-    TRANSCRIPTION              = args.TRANSCRIPTION
-    TRANSCRIPTION_PATH         = '{}_transcription.txt'.format(TRANSCRIPTION)
-    MODEL                      = args.MODEL
-    EMB_MODEL                  = args.EMB_MODEL
-
-    podcast_url_video_df = get_podcast_data(TRANSCRIPTION_PATH)
 
-    r = requests.get("https://raw.githubusercontent.com/AlbertoUAH/Castena/main/media/castena-animated-icon.gif", stream=True)
+@st.cache
+def setup_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((100, 100))
+    icon = icon.resize((img_size, img_size))
     st.sidebar.image(icon)
     video_option = st.sidebar.selectbox(
         "Seleccione el podcast",
@@ -72,11 +58,30 @@ def main():
     youtube_video_url   = list(podcast_url_video_df[podcast_url_video_df['podcast_name'].str.contains(video_option_joined)]['youtube_video_url'])[0].replace("\'", "")
 
     # -- 4. Setup request for system prompt
-    f = urllib.request.urlopen(DEFAULT_SYSTEM_PROMPT_LINK)
+    f = urllib.request.urlopen(default_system_prompt_link)
     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)
+    translator, nlp, retriever = utils.setup_app(video_option_joined_path, emb_model, model, logger)
+    return translator, nlp, retriever, video_option, video_option_joined_path
+    
+def main():
+    args = get_args()
+    B_INST, E_INST = "[INST]", "[/INST]"
+    B_SYS, E_SYS   = "<<SYS>>\n", "\n<</SYS>>\n\n"
+
+    PODCAST_URL_VIDEO_PATH     = args.PODCAST_URL_VIDEO_PATH
+    DEFAULT_SYSTEM_PROMPT_LINK = args.DEFAULT_SYSTEM_PROMPT_LINK
+    TRANSCRIPTION              = args.TRANSCRIPTION
+    TRANSCRIPTION_PATH         = '{}_transcription.txt'.format(TRANSCRIPTION)
+    MODEL                      = args.MODEL
+    EMB_MODEL                  = args.EMB_MODEL
+
+    podcast_url_video_df = get_podcast_data(TRANSCRIPTION_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)
 
 
     # -- 6. Setup prompt template + llm chain