Update streamlit_utils.py

streamlit_utils.py

@@ -50,21 +50,18 @@ def render_query():
     )
 
 
-@st.cache_data()
+@st.cache_data(persist=True)
 def load_model():
     model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
-    tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
     model = TFAutoModel.from_pretrained(model_ckpt, from_pt=True)
 
-    return tokenizer, model
+    return model
 
 
-@st.cache_data()
 def load_peft_model():
     peft_model_base = AutoModelForSeq2SeqLM.from_pretrained(
         "google/flan-t5-small", torch_dtype=torch.bfloat16
     )
-    peft_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
 
     peft_model = PeftModel.from_pretrained(
         peft_model_base,
@@ -72,24 +69,22 @@ def load_peft_model():
         torch_dtype=torch.bfloat16,
         is_trainable=False,
     )
-    return peft_model, peft_tokenizer
+    return peft_model
 
 
-@st.cache_data()
+@st.cache_data(persist=True)
 def load_faiss_dataset():
     faiss_dataset = load_dataset(
         "vishnupriyavr/wiki-movie-plots-with-summaries-faiss-embeddings",
         split="train",
     )
-    faiss_dataset.set_format("pandas")
-    df = faiss_dataset[:]
-    plots_dataset = Dataset.from_pandas(df)
-    plots_dataset.add_faiss_index(column="embeddings")
-    return plots_dataset
+    return faiss_dataset
 
 
 def get_embeddings(text_list):
-    tokenizer, model = load_model()
+    model = load_model()
+    model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
+    tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
     encoded_input = tokenizer(
         text_list, padding=True, truncation=True, return_tensors="tf"
     )
@@ -105,7 +100,11 @@ def cls_pooling(model_output):
 def search_movie(user_query, limit):
     question_embedding = get_embeddings([user_query]).numpy()
 
-    plots_dataset = load_faiss_dataset()
+    faiss_dataset = load_faiss_dataset()
+    faiss_dataset.set_format("pandas")
+    df = faiss_dataset[:]
+    plots_dataset = Dataset.from_pandas(df)
+    plots_dataset.add_faiss_index(column="embeddings")
     scores, samples = plots_dataset.get_nearest_examples(
         "embeddings", question_embedding, k=limit
     )
@@ -129,7 +128,8 @@ def search_movie(user_query, limit):
 
 
 def summarized_plot(sample_df, limit):
-    peft_model, peft_tokenizer = load_peft_model()
+    peft_model = load_peft_model()
+    peft_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
     peft_model_text_output_list = []
 
     for i in range(limit):
@@ -169,4 +169,4 @@ def aggregate(items):
         result["title"] = group[0]["title"]  # get titl from first item
         result["text"] = "\n\n".join([item["text"] for item in group])
         results.append(result)
-    return results
+    return results