main.py

import logging
import os

import streamlit as st
import pandas as pd
from db_setup import DBTableSetup
from query_execution import QueryExecution
from text_to_text_semantic_search import create_and_save_faiss_index_with_ids
from constants import *

# logging.basicConfig(level=logging.DEBUG)

if __name__ == '__main__':

  data_files = ["data/dataset.csv"]

  # UI Elements
  st.set_page_config(layout="wide", page_title="Semantic SQL")
  st.title("Semantic SQL")

  col1, col2 = st.columns([5, 5])

  selectedData = st.sidebar.selectbox('Choose Data', data_files)
  if st.sidebar.button("Setup DB"):
    base_table_setup = DBTableSetup(f"{DB_URL}/{DB_NAME}")
    d = pd.read_csv(selectedData)
    base_table_setup.insert_blob_data("songs", d, ["id"])
    create_and_save_faiss_index_with_ids("data/embeddings-instances.pt", "data/textIds-instances.pt",
                                         index_path="data/index", ids_save_path="data/index_ids.npy")

  if st.sidebar.button("Reset DB", key="db_reset"):
    # delete everything in database
    os.remove(DB_NAME)
    os.remove("data/index")
    os.remove("data/index_ids.npy")
    st.session_state["curr_executing"] = False

  # if st.sidebar.button("New Query", key="new_query"):
  #   # delete everything in database
  #   st.session_state["curr_executing"] = False
  #   st.session_state["results_ready"] = False

  if "curr_executing" not in st.session_state:
    st.session_state["curr_executing"] = False

  if "results_ready" not in st.session_state:
    st.session_state["results_ready"] = False

  if "curr_query" not in st.session_state:
    #TODO: Query badalni hai
    st.session_state["curr_query"] = "SELECT DISTINCT id from songs SEMANTIC 'id = 0'"#"""SELECT DISTINCT id from songs SEMANTIC 'title = "oh " '"""

  with col1:
    col1.header("Query the Data")
    input_query = st.text_area('Type SQL query here', st.session_state["curr_query"], 300)
    if st.button("Execute query"):
      st.session_state.curr_executing = True
      st.session_state.results_ready = False
      st.session_state["curr_query"] = input_query
      st.session_state["query_execution"] = QueryExecution(input_query)
      if st.session_state["query_execution"].execution_complete:
        st.session_state["results_ready"] = True

  with col2:
    col2.header("Query Results")
  if st.session_state.results_ready:
    with col2:
      results = st.session_state["query_execution"].get_results()
      st.write(results.head(10))