Update src/streamlit_app.py

src/streamlit_app.py

@@ -68,16 +68,24 @@ def show_img(img_id, score=None, col=None):
             st.error(f"Error processing image: {e}")
 
 def search_faiss(model, processor, index, id_map, prompt, top_k=5, device='cpu'):
+    st.write(f"Running FAISS search for prompt: '{prompt}' with top_k={top_k}")
     inputs = processor(text=[prompt], return_tensors='pt', padding=True).to(device)
+    st.write("Prompt processed by tokenizer.")
+    
     with torch.no_grad():
         txt_emb = model.get_text_features(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
         txt_emb = txt_emb / txt_emb.norm(p=2, dim=-1, keepdim=True)
+    st.write("Text embedding computed.")
+    
     q = txt_emb.cpu().numpy().astype('float32')
 
     D, I = index.search(q, top_k)
+
+    st.write("FAISS search completed.")
     return [(id_map[i], float(D[0][j])) for j, i in enumerate(I[0])]
 
 def running(prompt, top_k=5):
+    st.write("Starting image retrieval...")
     results = search_faiss(
         model, processor,
         index, id_map,
@@ -175,5 +183,5 @@ st.subheader("Example Usage: Enter a Prompt to Retrieve Related Images")
 prompt_input = st.text_input("Enter a prompt", "a red Apparel dress")
 top_k_input = st.number_input("Enter the number of results (top_k)", min_value=1, max_value=10, value=5)
 
-if st.button("Find top 5 related images:"):
+if st.button(f"Find top {top_k_input} related images"):
     running(prompt_input, top_k_input)