Update app.py

app.py

@@ -76,16 +76,20 @@ def semantic_search(q: str, top_k: int=100):
     D, I = faiss_index.search(emb, top_k)
     return [(etf_list[i], float(D[0][j])) for j,i in enumerate(I[0])]
 
-def ensemble_ticker_extraction(q: str):
+
+# Ensemble function: union of both models' predictions
+def ensemble_ticker_extraction(query):
     preds = set()
+
     for tok, mdl in ((tok1,m1),(tok2,m2)):
-        enc = tok(q, return_tensors="pt")
+        enc = tok(query, return_tensors="pt")
         with torch.no_grad():
             logits = mdl(**enc).logits
-        ids  = logits.argmax(dim=-1)[0].tolist()
-        toks = tok.convert_ids_to_tokens(enc["input_ids"][0])
-        labs = [mdl.config.id2label[i] for i in ids]
-        preds |= modules.extract_valid_tickers(toks, labs, tok, valid_ticker_set)
+        pred_ids = logits.argmax(dim=-1)[0].tolist()
+        tokens   = tok.convert_ids_to_tokens(enc["input_ids"][0])
+        labels   = [mdl.config.id2label[i] for i in pred_ids]
+        preds.update(modules.extract_valid_tickers(tokens, labels, tok, valid_ticker_set))
+
     return preds
 
 # ─── UI HELPERS ─────────────────────────────────────────────────────────────
@@ -169,12 +173,16 @@ def display_chat_history(task: str):
 def process_query(task: str, query: str):
     top_k, top_n = 100, 30
     if task=="search_etf":
+        """Process a query by calling your LLM and appending results."""
+        # st.session_state["all_chat_history"].append(app_helper_utilities.form_d_chat_history(result_id=str(uuid.uuid4()), response=query_text, task='user_query'))
+        # st.session_state["explore_conversation"].append({"role": "user", "content": query_text})
+        st.chat_message("user").write(query)
         with st.spinner("Searching ETFs..."):
             fetched = semantic_search(query, top_k)
             df_out  = modules.get_etf_recommendations_from_list(
                 fetched, df_etf, top_n
             )
-        msg = f"{len(df_out)} ETFs found."
+        msg = f"{len(df_out)} {fetched} ETFs found."
         st.session_state[f"all_chat_history_{task}"].append(
             modules.form_d_chat_history(str(uuid.uuid4()), msg, task, df=df_out)
         )