Add application file

app.py

@@ -0,0 +1,517 @@
+import time
+start_time = time.time()
+from pathlib import Path
+from typing import Tuple
+import pandas as pd
+import gradio as gr
+import json
+
+import duckdb
+from sentence_transformers import SentenceTransformer
+from datasets import load_dataset
+
+USE_DOTENV = False
+
+ROOT = Path(__file__).parent
+
+JSON_PATH = ROOT / "json"
+# DATASET_PATH = ROOT / "pkl" / "app_dataset.pkl"
+DOTENV_PATH = ROOT.parent.parent / "apis" / ".env"
+# DUCKDB_PATH = ROOT / "db" / "sss_vectordb.duckdb"
+
+from src import front_dataset_handler as fdh, app_utils as utils, semantic_search as ss, env_options
+tokens = env_options.check_env(use_dotenv=USE_DOTENV, dotenv_path=DOTENV_PATH, env_tokens = ["HF_TOKEN"])
+print(f"Libraries loaded. {time.time() - start_time:.2f} seconds.")
+# Carga de modelo de embeddings y conexión a DuckDB
+emb_model = SentenceTransformer("FinLang/finance-embeddings-investopedia", token = tokens.get("HF_TOKEN"))
+# con = duckdb.connect(DUCKDB_PATH)
+print(f"Model loaded. {time.time() - start_time:.2f} seconds.")
+#### CONEXIÓN DUCKDB A HUGGING FACE HUB ####
+print("Initializing DuckDB connection...")
+con = duckdb.connect()
+hf_token = tokens.get("HF_TOKEN")
+##################################
+masked_hf_token = hf_token[:4] + "*" * (len(hf_token) - 8) + hf_token[-4:]
+print(f"Using Hugging Face token: {masked_hf_token}")
+##################################
+
+hf_token = tokens.get("HF_TOKEN")
+masked_hf_token = hf_token[:4] + "*" * (len(hf_token) - 8) + hf_token[-4:]
+'''
+create_secret_query = f"""
+    INSTALL httpfs;
+    LOAD httpfs;
+    CREATE PERSISTENT SECRET hf_token (
+        TYPE huggingface,
+        TOKEN '{hf_token}'
+    );
+    """
+'''
+# con.sql(create_secret_query)
+# print(con.sql("SELECT * FROM duckdb_secrets()").fetchdf())
+dataset_name = "reddgr/swift-stock-screener"
+# con.sql(query="INSTALL vss; LOAD vss;")
+
+create_secret_query = f"""
+        INSTALL httpfs;
+        LOAD httpfs;
+        CREATE PERSISTENT SECRET hf_token (
+            TYPE huggingface,
+            TOKEN '{hf_token}'
+        );
+        """
+con.sql(create_secret_query)
+print(con.sql("SELECT * FROM duckdb_secrets()").fetchdf().iloc[0,-2])
+print(con.sql("SELECT * FROM duckdb_secrets()").fetchdf().iloc[0,-1])
+print(con.sql("SELECT * FROM duckdb_secrets()").fetchdf())
+
+create_table_query = f"""
+        INSTALL vss;
+        LOAD vss;
+        SET hnsw_enable_experimental_persistence = true;
+        CREATE TABLE vector_table AS 
+        SELECT *, embeddings::float[{emb_model.get_sentence_embedding_dimension()}] as embeddings_float 
+        FROM 'hf://datasets/{dataset_name}/data/train-00000-of-00001.parquet';
+        """
+
+con.sql(create_table_query)
+
+print("Indexing data for vector search...")
+create_index_query = f"""
+        CREATE INDEX sss_hnsw_index ON vector_table USING HNSW (embeddings_float) WITH (metric = 'cosine');
+        """
+con.sql(create_index_query)
+
+# print(con.sql("SELECT * FROM duckdb_secrets()").fetchdf())
+print(f"Created search index. {time.time() - start_time:.2f} seconds.")
+########################################
+
+# ESTADO GLOBAL
+last_result_df: pd.DataFrame = pd.DataFrame()
+
+######################
+last_search_type: str = ""
+last_search_query: str = ""
+# last_filtros_values: Tuple = ()
+last_column_filters: list[tuple[str, str]] = []
+last_sort_col_label: str = ""
+last_sort_dir: str = ""
+#######################
+
+# ---------------------------------------------------------------------------
+# CONFIG --------------------------------------------------------------------
+# ---------------------------------------------------------------------------
+app_dataset = load_dataset("reddgr/swift-stock-screener", split="train", token = tokens.get("HF_TOKEN")).to_pandas()
+
+# dh_app = fdh.FrontDatasetHandler(app_dataset=pd.read_pickle(DATASET_PATH))
+dh_app = fdh.FrontDatasetHandler(app_dataset=app_dataset)
+maestro = dh_app.app_dataset[dh_app.app_dataset['quoteType']=='EQUITY'].copy()
+maestro_etf = dh_app.app_dataset[dh_app.app_dataset['quoteType']=='ETF'].copy()
+
+with open(JSON_PATH / "app_column_config.json", "r") as f:
+    variables_busq_norm = json.load(f)["variables_busq_norm"]
+
+with open(JSON_PATH / "app_column_config.json", "r") as f:
+    caracteristicas = json.load(f)["cols_tabla_equity"]
+
+with open(JSON_PATH / "app_column_config.json", "r") as f:
+    caracteristicas_etf = json.load(f)["cols_tabla_etfs"]
+
+with open(JSON_PATH / "cat_cols.json", "r") as f:
+    cat_cols = json.load(f)["cat_cols"]
+
+with open(JSON_PATH / "col_names_map.json", "r") as f:
+    rename_columns = json.load(f)["col_names_map"]
+
+with open(JSON_PATH / "gamma_params.json", "r") as f:
+    gamma_params = json.load(f)
+
+with open(JSON_PATH / "semantic_search_params.json", "r") as f:
+    semantic_search_params = json.load(f)["semantic_search_params"]
+
+# Columnas a estilizar en rojo si son negativas
+neg_display_cols = [rename_columns.get(c, c)
+                    for c in ("ret_365", "revenueGrowth")]
+
+# Parámetros de la función de distribución de distancias:
+shape, loc, scale = gamma_params["shape"], gamma_params["loc"], gamma_params["scale"]
+max_dist, precision_cdf = gamma_params["max_dist"], gamma_params["precision_cdf"]
+y_cdf, _ = dh_app.configura_distr_prob(shape, loc, scale, max_dist, precision_cdf)
+
+# Parámetros de la de búsqueda VSS:
+k = semantic_search_params["k"]
+brevity_penalty = semantic_search_params["brevity_penalty"]
+reward_for_literal = semantic_search_params["reward_for_literal"]
+partial_match_factor = semantic_search_params["partial_match_factor"]
+print(f"VSS params: k={k}, brevity_penalty={brevity_penalty}, reward_for_literal={reward_for_literal}, partial_match_factor={partial_match_factor}")
+
+filtros_keys = caracteristicas[2:]
+
+MAX_ROWS = 13000
+ROWS_PER_PAGE = 100
+
+# ---------------------------------------------------------------------------
+# FUNCIONES UI --------------------------------------------------------------
+# ---------------------------------------------------------------------------
+
+# Dejamos en este módulo (en lugar de app_utils) funciones específicas de gestión de la interfaz
+
+def _paginate(df: pd.DataFrame, page: int, per_page: int = ROWS_PER_PAGE) -> Tuple[pd.DataFrame, str]:
+    total_pages = max(1, (len(df) + per_page - 1) // per_page)
+    page        = max(1, min(page, total_pages))
+    slice_df    = df.iloc[(page-1)*per_page : (page-1)*per_page + per_page]
+    slice_df = utils.styler_negative_red(slice_df, cols=neg_display_cols)
+    return slice_df, f"Page {page} of {total_pages}"
+
+
+def search_dynamic(ticker: str, page: int, *filtros_values) -> Tuple[pd.DataFrame, str]:
+    global last_result_df
+
+    ticker = ticker.upper().strip()
+    if ticker == "":
+        last_result_df = pd.DataFrame()
+        return pd.DataFrame(), "Page 1 of 1"
+
+    filtros = dict(zip(filtros_keys, filtros_values))
+
+    neighbors_df = dh_app.vecinos_cercanos(
+        df=maestro,
+        variables_busq=variables_busq_norm,
+        caracteristicas=caracteristicas,
+        target_ticker=ticker,
+        y_cdf=y_cdf,
+        precision_cdf=precision_cdf,
+        max_dist=max_dist,
+        n_neighbors=len(maestro),
+        filtros=filtros,
+    )
+
+    if isinstance(neighbors_df, str):
+        last_result_df = pd.DataFrame()
+        return pd.DataFrame(), "Page 1 de 1"
+
+    neighbors_df.reset_index(inplace=True)
+    neighbors_df.drop(columns=["distance"], inplace=True)
+    # neighbors_df = format_results(neighbors_df)
+    neighbors_df = utils.format_results(neighbors_df, rename_columns)
+
+    last_result_df = neighbors_df.head(MAX_ROWS).copy()
+    return _paginate(last_result_df, page)
+
+
+def search_theme(theme: str, page: int, *filtros_values) -> Tuple[pd.DataFrame, str]:
+    global last_result_df
+    query = theme.strip()
+    if query == "":
+        last_result_df = pd.DataFrame()
+        return pd.DataFrame(), "Page 1 of 1"
+
+    # Llamada al algoritmo de búsqueda, que devuelve un dataframe con k activos:
+    result_df = ss.duckdb_vss_local(
+        model=emb_model,
+        duckdb_connection=con,
+        query=query,
+        k=k,
+        brevity_penalty=brevity_penalty,
+        reward_for_literal=reward_for_literal,
+        partial_match_factor=partial_match_factor,
+        table_name="vector_table",
+        embedding_column="embeddings"
+    )
+    theme_dist = result_df[['ticker', 'distance']].rename(columns={'distance': 'Search dist.'})
+    # Cruzamos el dataframe de distancias con el maestro y mantenemos las columnas originales: 
+    clean_feats = [c for c in caracteristicas if c != 'ticker']
+    # indexamos por ticker para cruzar las tablas:
+    maestro_subset = maestro.set_index('ticker')[clean_feats]
+    merged = theme_dist.set_index('ticker').join(maestro_subset, how='inner').reset_index()
+    # Reordenamos las columnas y añadimos la distancia:
+    ordered_cols = ['ticker'] + clean_feats + ['Search dist.']
+    merged = merged[ordered_cols]
+    # Ajustamos los formatos de las columnas:
+    formatted = utils.format_results(merged, rename_columns)
+    last_result_df = formatted.head(MAX_ROWS).copy()
+    return _paginate(last_result_df, page)
+
+
+def _compose_summary() -> str:
+    parts = []
+    if last_search_type == "theme":
+        parts.append(f"Theme search for '{last_search_query}'")
+    elif last_search_type == "ticker":
+        parts.append(f"Ticker search for '{last_search_query}'")
+    if last_column_filters:
+        fstr = ", ".join(f"{col} = '{val}'" for col, val in last_column_filters)
+        parts.append(f"Filters: {fstr}")
+    if last_sort_col_label:
+        parts.append(f"Sorted by: {last_sort_col_label} ({last_sort_dir})")
+    return ". ".join(parts)
+
+def search_all(theme: str, ticker: str, page: int) -> tuple[pd.DataFrame,str,str,str,str]:
+    global last_search_type, last_search_query, last_column_filters
+    last_column_filters.clear()
+
+    if theme.strip():
+        last_search_type, last_search_query = "theme", theme.strip()
+        df, label = search_theme(theme, page)
+        # new_ticker, new_theme = "", theme.strip()
+        new_ticker, new_theme = "", "" # limpia las cajas de búsqueda
+
+    elif ticker.strip():
+        last_search_type, last_search_query = "ticker", ticker.strip().upper()
+        df, label = search_dynamic(ticker, page)
+        # new_ticker, new_theme = last_search_query, ""
+        new_ticker, new_theme = "", ""
+
+    else:
+        df, label = _paginate(last_result_df, page)
+        new_ticker, new_theme = "", ""
+
+    summary = _compose_summary()
+    return df, label, new_ticker, new_theme, summary
+
+def page_change(theme: str, ticker: str, page: int) -> tuple[pd.DataFrame,str,str,str,str]:
+    return search_all(theme, ticker, page)
+
+
+# ---------------------------------------------------------------------------
+# SORTING -------------------------------------------------------------------
+# ---------------------------------------------------------------------------
+
+def apply_sort(col_label: str, direction: str) -> tuple[pd.DataFrame, str, int, str]:
+    global last_sort_col_label, last_sort_dir, last_search_type, last_search_query, last_column_filters, last_result_df
+
+    # record selection and clear previous state
+    last_sort_col_label, last_sort_dir = col_label or "", direction or ""
+    last_search_type = last_search_query = ""
+    last_column_filters.clear()
+
+    # reload raw data
+    df_raw = maestro[caracteristicas].head(MAX_ROWS).copy()
+
+    # sort on original data column if specified
+    if col_label:
+        # reverse lookup original column key
+        inv_map = {v: k for k, v in rename_columns.items()}
+        orig_col = inv_map.get(col_label, col_label)
+        asc = (direction == "Ascending")
+        df_raw = df_raw.sort_values(
+            by=orig_col,
+            ascending=asc,
+            na_position='last'
+        ).reset_index(drop=True)
+
+    # apply existing formatting helpers
+    df_formatted = utils.format_results(df_raw, rename_columns)
+
+    # update global and paginate
+    last_result_df = df_formatted.copy()
+    slice_df, label = _paginate(last_result_df, 1)
+    summary = f"Sorted by: {col_label} ({direction})" if col_label else ""
+    return slice_df, label, 1, summary
+
+
+
+def reset_initial() -> tuple[pd.DataFrame,str,int,str,str,str]:
+    global last_search_type, last_search_query, last_column_filters, last_sort_col_label, last_sort_dir, last_result_df
+    last_search_type = last_search_query = ""
+    last_column_filters.clear()
+    last_sort_col_label = last_sort_dir = ""
+    last_result_df = utils.format_results(maestro[caracteristicas].head(MAX_ROWS).copy(), rename_columns)
+    slice_df, label = _paginate(last_result_df, 1)
+    default_sort = rename_columns.get("marketCap","marketCap")
+    return slice_df, label, 1, "", "", default_sort, ""
+
+
+# ---------------------------------------------------------------------------
+# DATOS INICIALES -----------------------------------------------------------
+# ---------------------------------------------------------------------------
+
+last_result_df = utils.format_results(maestro[caracteristicas].head(MAX_ROWS).copy(), rename_columns)
+_initial_slice, _initial_label = _paginate(last_result_df, 1)
+
+# ---------------------------------------------------------------------------
+# UI ------------------------------------------------------------------------
+# ---------------------------------------------------------------------------
+
+def _load_html(name: str) -> str:
+    return (ROOT / "html" / name).read_text(encoding="utf-8")
+
+html_front_layout = _load_html("front_layout.html")
+
+with gr.Blocks(title="Swift Stock Screener, by Reddgr") as front:
+    gr.HTML(html_front_layout)
+
+    # ---------------------- TOP INPUT -------------------------------------
+    with gr.Row(equal_height=True):
+        theme_input = gr.Textbox(show_label=False, placeholder="Search a theme. i.e. 'lithium'", scale=2)
+        ticker_input = gr.Textbox(show_label=False, placeholder="Enter a ticker symbol", scale=1)   
+        buscar_button = gr.Button("Search")
+        gr.HTML("<div></div>")
+        reset_button = gr.Button("Reset", elem_classes="small-btn")
+        # gr.HTML("<div></div>")
+        random_button = gr.Button("Random ticker", elem_classes="small-btn")
+
+    # ---------------------- SEARCH SUMMARY ------------------------
+    summary_display = gr.Markdown("", elem_classes="search-spec")
+
+    # ---------------------- DATAFRAME & PAGINATION ------------------------
+
+    output_df = gr.Dataframe(
+        value=_initial_slice,
+        interactive=False,
+        elem_classes="clickable-columns",
+        # max_height=500
+    )
+
+
+    # ---------------------- PAGINATION AND SORT CONTROLS ---------------------
+    with gr.Row():
+        btn_prev = gr.Button("Previous", elem_classes="small-btn")
+        pagination_label = gr.Markdown(_initial_label)
+        btn_next = gr.Button("Next", elem_classes="small-btn")
+        gr.Markdown("&nbsp;" * 20)
+        # merged sort controls on right
+        sort_col = gr.Dropdown(
+            choices=[rename_columns.get(c, c) for c in caracteristicas],
+            value=None,
+            label="Reset and sort by:",
+            allow_custom_value=False,
+            scale=2,
+        )
+        sort_dir = gr.Radio(
+            choices=["Ascending", "Descending"],
+            value="Descending",
+            label="",
+            scale=1,
+        )
+
+    page_state = gr.State(1)
+
+    # ---------------------- EXCLUSION FILTER TOGGLES --------------------------------
+    # De momento excluimos esta funcionalidad, al menos en la tabla de acciones, 
+    # por la complejidad que añade (es una herencia del buscador de fondos de inversión)
+    # Potencial mejora para cuando incorporemos la tabla de ETFs
+    '''
+    with gr.Row():
+        toggle_components = [
+            gr.Checkbox(value=True, label=rename_columns.get(k, k)) for k in filtros_keys
+        ]
+    '''
+
+    # ---------------------- HELPERS ---------------------------------------
+    def reset_page():
+        return 1
+
+    def prev_page(p):
+        return max(p - 1, 1)
+
+    def next_page(p):
+        return p + 1
+
+    def search_inputs():
+        return [theme_input, ticker_input, page_state]
+
+    def random_action() -> tuple[str,int,str]:
+        return utils.random_ticker(maestro), 1, ""
+
+    # ---------------------- BINDINGS --------------------------------------
+    # search_dynamic -> search_all
+    inputs = [theme_input, ticker_input, page_state]
+
+    buscar_button.click(
+        search_all,
+        inputs=inputs,
+        outputs=[output_df, pagination_label, ticker_input, theme_input, summary_display]
+    )
+
+    ticker_input.submit(
+        reset_page, None, page_state
+    ).then(
+        search_all,
+        inputs=inputs,
+        outputs=[output_df, pagination_label, ticker_input, theme_input, summary_display]
+    )
+
+    theme_input.submit(
+        reset_page, None, page_state
+    ).then(
+        search_all,
+        inputs=inputs,
+        outputs=[output_df, pagination_label, ticker_input, theme_input, summary_display]
+    )
+
+    random_button.click(
+        random_action,
+        None,
+        [ticker_input, page_state, theme_input]
+    ).then(
+        search_all,
+        inputs=inputs,
+        outputs=[output_df, pagination_label, ticker_input, theme_input, summary_display]
+    )
+
+    reset_button.click(
+        reset_initial,
+        None,
+        [output_df, pagination_label, page_state, ticker_input, theme_input, sort_col, summary_display]
+    )
+
+    btn_prev.click(
+        prev_page, page_state, page_state
+    ).then(
+        page_change,
+        inputs=inputs,
+        outputs=[output_df, pagination_label, ticker_input, theme_input, summary_display]
+    )
+
+    btn_next.click(
+        next_page, page_state, page_state
+    ).then(
+        page_change,
+        inputs=inputs,
+        outputs=[output_df, pagination_label, ticker_input, theme_input, summary_display]
+    )
+
+    sort_col.change(
+        apply_sort,
+        inputs=[sort_col, sort_dir],
+        outputs=[output_df, pagination_label, page_state, summary_display]
+    )
+
+    sort_dir.change(
+        apply_sort,
+        inputs=[sort_col, sort_dir],
+        outputs=[output_df, pagination_label, page_state, summary_display]
+    )
+
+    # ---------------------- FILTERS BY COLUMN ------------------ #
+    filterable_columns = [rename_columns.get(c, c) for c in cat_cols]
+
+
+    def filter_by_column(evt: gr.SelectData) -> tuple[pd.DataFrame,str,int,str]:
+        global last_result_df, last_column_filters
+        if last_result_df.empty:
+            return pd.DataFrame(), "Page 1 of 1", 1, _compose_summary()
+
+        col = last_result_df.columns[evt.index[1]]
+        # print(f"DEBUG: resolving to column #{evt.index[1]} → '{col}'")
+        val = evt.value
+        last_column_filters.append((col, val))
+        filtered = last_result_df[last_result_df[col] == val]
+        last_result_df = filtered.copy()
+        slice_df, label = _paginate(last_result_df, 1)
+        summary = _compose_summary()
+        return slice_df, label, 1, summary
+
+
+    output_df.select(
+        filter_by_column,
+        outputs=[output_df, pagination_label, page_state, summary_display]
+    )
+
+# ---------------------------------------------------------------------------
+# LAUNCH --------------------------------------------------------------------
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    front.launch()