Quick bug fixes, adjustments, cosmetics #44

Author: pramitchoudhary

Makefile

@@ -10,7 +10,6 @@ setup: download_demo_data ## Setup
 	./.sidekickvenv/bin/python3 -m pip install --upgrade pip
 	./.sidekickvenv/bin/python3 -m pip install wheel
 	./.sidekickvenv/bin/python3 -m pip install -r requirements.txt
-	mkdir -p ./db/sqlite
 	mkdir -p ./examples/demo/
 
 download_models:

about.md

@@ -4,7 +4,7 @@
 
 **Actively Being Maintained:** Yes (Demo release: _In active RnD_)
 
-**Last Updated:** October, 2023
+**Last Updated:** November, 2023
 
 **Allows uploading and using new model and data:** Yes
 

app.toml

@@ -4,11 +4,11 @@ title = "SQL-Sidekick"
 description = "QnA with tabular data using NLQ"
 LongDescription = "about.md"
 Tags = ["DATA_SCIENCE", "MACHINE_LEARNING", "NLP"]
-Version = "0.1.3"
+Version = "0.1.4"
 
 [Runtime]
 MemoryLimit = "64Gi"
-MemoryReservation = "16Gi"
+MemoryReservation = "64Gi"
 module = "start"
 VolumeMount = "/meta_data"
 VolumeSize = "100Gi"

sidekick/db_config.py

@@ -39,6 +39,7 @@ def __init__(
         self.base_path = base_path
         self.column_names = []
         if dialect == "sqlite":
+            logger.debug(f"Creating SQLite DB: sqlite:///{base_path}/db/sqlite/{db_name}.db")
             self._url = f"sqlite:///{base_path}/db/sqlite/{db_name}.db"
         else:
             self._url = f"{self.dialect}://{self.user_name}:{self.password}@{self.hostname}:{self.port}/"
@@ -181,9 +182,10 @@ def add_samples(self, data_csv_path=None):
             # Fetch the number of rows from the table
             sample_query = f"SELECT COUNT(*) AS ROWS FROM {self.table_name} LIMIT 1"
             num_rows = pd.read_sql_query(sample_query, engine)
-            logger.info(f"Number of rows inserted: {num_rows.values[0][0]}")
+            res = num_rows.values[0][0]
+            logger.info(f"Number of rows inserted: {res}")
             engine.dispose()
-            return num_rows, None
+            return res, None
         except SQLAlchemyError as sqla_error:
             logger.debug("SQLAlchemy error:", sqla_error)
             return None, sqla_error

sidekick/prompter.py

@@ -28,12 +28,14 @@
 )
 
 # Load the config file and initialize required paths
-base_path = (Path(__file__).parent / "../").resolve()
-env_settings = toml.load(f"{base_path}/sidekick/configs/env.toml")
+app_base_path = (Path(__file__).parent / "../").resolve()
+# Below check is to handle the case when the app is running on the h2o.ai cloud or locally
+base_path = app_base_path if os.path.isdir("./.sidekickvenv/bin/") else "/meta_data"
+env_settings = toml.load(f"{app_base_path}/sidekick/configs/env.toml")
 db_dialect = env_settings["DB-DIALECT"]["DB_TYPE"]
 model_name = env_settings["MODEL_INFO"]["MODEL_NAME"]
 os.environ["TOKENIZERS_PARALLELISM"] = "False"
-__version__ = "0.0.4"
+__version__ = "0.1.4"
 
 
 def color(fore="", back="", text=None):
@@ -189,7 +191,7 @@ def db_setup_api(
         # env_settings["TABLE_INFO"]["TABLE_SAMPLES_PATH"] = table_samples_path
 
         # Update settings file for future use.
-        f = open(f"{base_path}/sidekick/configs/env.toml", "w")
+        f = open(f"{app_base_path}/sidekick/configs/env.toml", "w")
         toml.dump(env_settings, f)
         f.close()
         path = f"{base_path}/var/lib/tmp/data"
@@ -391,7 +393,7 @@ def query_api(
             env_settings["MODEL_INFO"]["OPENAI_API_KEY"] = api_key
 
             # Update settings file for future use.
-            f = open(f"{base_path}/sidekick/configs/env.toml", "w")
+            f = open(f"{app_base_path}/sidekick/configs/env.toml", "w")
             toml.dump(env_settings, f)
             f.close()
         openai.api_key = api_key

sidekick/query.py

@@ -75,6 +75,7 @@ def __new__(
         if cls._instance is None or (cls._instance and not cls._instance.models.get(model_name, None)):
             if cls._instance is None:
                 cls._instance = super().__new__(cls)
+                cls._instance.current_temps = {}
             cls._instance.models, cls._instance.tokenizers = load_causal_lm_model(
                 model_name,
                 cache_path=f"{job_path}/models/",
@@ -84,7 +85,7 @@ def __new__(
             )
             cls._instance.model_name = "h2ogpt-sql-sqlcoder2" if not model_name else model_name
             model_embed_path = f"{job_path}/models/sentence_transformers"
-            cls._instance.models[cls._instance.model_name].current_temperature = 0.5
+            cls._instance.current_temps[cls._instance.model_name] = 0.5
             device = "cuda" if torch.cuda.is_available() else "cpu" if device == "auto" else device
             cls._instance.similarity_model = load_embedding_model(model_path=model_embed_path, device=device)
         return cls._instance
@@ -479,6 +480,7 @@ def generate_sql(
                 tokenizer = self.tokenizers[model_name]
                 inputs = tokenizer([query], return_tensors="pt")
                 model = self.models[model_name]
+                current_temperature = self.current_temps.get(model_name, 0.5)
                 input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]
                 logger.info(f"Context length: {input_length}")
 
@@ -512,19 +514,25 @@ def generate_sql(
 
                 possible_temp_gt_5 = [0.6, 0.75, 0.8, 0.9, 1.0]
                 possible_temp_lt_5 = [0.1, 0.2, 0.3, 0.4]
-                random_temperature = model.current_temperature
                 random_seed = random.randint(0, 50)
                 torch.manual_seed(random_seed)
-                if model.current_temperature >= 0.5:
+
+                if current_temperature >= 0.5:
                     random_temperature = np.random.choice(possible_temp_lt_5, 1)[0]
                 else:
                     random_temperature = np.random.choice(possible_temp_gt_5, 1)[0]
+                import pdb
+
+                pdb.set_trace()
                 if not self.is_regenerate_with_options and not self.is_regenerate:
                     # Greedy decoding
+                    # Reset temperature to 0.5
+                    current_temperature = 0.5
+                    logger.debug(f"Generation with default temperature : {current_temperature}")
                     output = model.generate(
                         **inputs.to(device_type),
                         max_new_tokens=512,
-                        temperature=0.5,
+                        temperature=current_temperature,
                         output_scores=True,
                         do_sample=True,
                         return_dict_in_generate=True,
@@ -544,7 +552,8 @@ def generate_sql(
                         return_dict_in_generate=True,
                     )
                     generated_tokens = output.sequences[:, input_length:][0]
-                    model.current_temperature = random_temperature
+                    self.current_temps[model_name] = random_temperature
+                    logger.debug(f"Temperature saved: {self.current_temps[model_name]}")
                 else:
                     logger.info("Regeneration with options requested on previous query ...")
                     # Diverse beam search decoding to explore more options

sidekick/utils.py

@@ -29,8 +29,8 @@
 }
 
 TASK_CHOICE = {
-    "q_a": "Question/Answering",
-    "sqld": "SQL Debugging",
+    "q_a": "Ask Questions",
+    "sqld": "Debugging",
 }
 
 
@@ -60,6 +60,7 @@ def generate_sentence_embeddings(model_path: str, x, batch_size: int = 32, devic
 
 
 def load_embedding_model(model_path: str, device: str):
+    logger.debug(f"Loading embedding model from: {model_path}")
     model_name_path = glob.glob(f"{model_path}/models--BAAI--bge-base-en/snapshots/*/")[0]
 
     sentence_model = SentenceTransformer(model_name_path, cache_folder=model_path, device=device)
@@ -186,7 +187,7 @@ def save_query(
 
 
 def setup_dir(base_path: str):
-    dir_list = ["var/lib/tmp/data", "var/lib/tmp/jobs", "var/lib/tmp/.cache", "models/weights"]
+    dir_list = ["var/lib/tmp/data", "var/lib/tmp/jobs", "var/lib/tmp/.cache", "models", "db/sqlite"]
     for _dl in dir_list:
         p = Path(f"{base_path}/{_dl}")
         if not p.is_dir():
@@ -344,7 +345,7 @@ def _load_llm(model_type: str, device_index: int = 0, load_in_4bit=True):
             _load_in_8bit = load_in_8bit
             model_name = model_type
             logger.info(f"Loading model: {model_name} on device id: {device_index}")
-
+            logger.debug(f"Model cache: {cache_path}")
             # 22GB (Least requirement on GPU) is a magic number for the current model size.
             if off_load and re_generate and total_memory < 22:
                 # To prevent the system from crashing in-case memory runs low.

start.py

@@ -6,18 +6,25 @@
 
 from huggingface_hub import snapshot_download
 
-print(f"Download model...")
-base_path = (Path(__file__).parent).resolve()
-
-MODEL_CHOICE_MAP = {
-    "h2ogpt-sql-sqlcoder2": "defog/sqlcoder2",
-    "h2ogpt-sql-nsql-llama-2-7B": "NumbersStation/nsql-llama-2-7B",
-}
-
-for _m in MODEL_CHOICE_MAP.values():
-    print(f"Downloading {_m}...", flush=True)
-    snapshot_download(repo_id=_m, cache_dir=f"{base_path}/models/")
-    time.sleep(3)
+
+def setup_dir(base_path: str):
+    dir_list = ["var/lib/tmp/data", "var/lib/tmp/jobs", "var/lib/tmp/.cache", "models", "db/sqlite"]
+    for _dl in dir_list:
+        p = Path(f"{base_path}/{_dl}")
+        if not p.is_dir():
+            p.mkdir(parents=True, exist_ok=True)
+
+
+print(f"Download models...")
+base_path = (Path(__file__).parent).resolve() if os.path.isdir("./.sidekickvenv/bin/") else "/meta_data"
+setup_dir(base_path)
+
+# Model 1:
+print(f"Download model 1...")
+snapshot_download(repo_id="NumbersStation/nsql-llama-2-7B", cache_dir=f"{base_path}/models/")
+# Model 2:
+print(f"Download model 2...")
+snapshot_download(repo_id="defog/sqlcoder2", cache_dir=f"{base_path}/models/")
 
 print(f"Download embedding model...")
 snapshot_download(repo_id="BAAI/bge-base-en", cache_dir=f"{base_path}/models/sentence_transformers/")

ui/app.py

@@ -15,8 +15,10 @@
 from sidekick.utils import TASK_CHOICE, get_table_keys, save_query, setup_dir, update_tables
 
 # Load the config file and initialize required paths
-base_path = (Path(__file__).parent / "../").resolve()
-env_settings = toml.load(f"{base_path}/ui/app_config.toml")
+app_base_path = (Path(__file__).parent / "../").resolve()
+env_settings = toml.load(f"{app_base_path}/ui/app_config.toml")
+# Below check is to handle the case when the app is running on the h2o.ai cloud or locally
+base_path = app_base_path if os.path.isdir("./.sidekickvenv/bin/") else "/meta_data"
 tmp_path = f"{base_path}/var/lib/tmp"
 
 ui_title = env_settings["WAVE_UI"]["TITLE"]
@@ -45,7 +47,7 @@ def initialize_models():
 
 
 async def user_variable(q: Q):
-    db_settings = toml.load(f"{base_path}/sidekick/configs/env.toml")
+    db_settings = toml.load(f"{app_base_path}/sidekick/configs/env.toml")
 
     q.user.db_dialect = db_settings["DB-DIALECT"]["DB_TYPE"]
     q.user.host_name = db_settings["LOCAL_DB_CONFIG"]["HOST_NAME"]
@@ -115,7 +117,7 @@ async def chat(q: Q):
     ]
     q.user.model_choice_dropdown = "h2ogpt-sql-sqlcoder2"
 
-    task_choices = [ui.choice("q_a", "Question/Answering"), ui.choice("sqld", "SQL Debugging")]
+    task_choices = [ui.choice("q_a", "Ask Questions"), ui.choice("sqld", "Debugging")]
     q.user.task_choice_dropdown = "q_a"
     add_card(
         q,
@@ -162,7 +164,7 @@ async def chat(q: Q):
             items=[
                 ui.dropdown(
                     name="task_dropdown",
-                    label="Task",
+                    label="Mode",
                     required=True,
                     choices=task_choices,
                     value=q.user.task_choice_dropdown if q.user.task_choice_dropdown else None,
@@ -250,7 +252,7 @@ async def chatbot(q: Q):
     if (
         f"Table {q.user.table_dropdown} selected" in q.args.chatbot
         or f"Model {q.user.model_choice_dropdown} selected" in q.args.chatbot
-        or f"Task {q.user.task_dropdown} selected" in q.args.chatbot
+        or f"{q.user.task_dropdown} mode selected" in q.args.chatbot
     ):
         return
 
@@ -417,13 +419,13 @@ async def fileupload(q: Q):
                 table_name=q.user.table_name,
             )
             logging.info(f"DB updates: \n {db_resp}")
-            q.args.n_rows = n_rows
             if "error" in str(db_resp).lower():
                 q.page["dataset"].error_upload_bar.visible = True
                 q.page["dataset"].error_bar.visible = False
                 q.page["dataset"].progress_bar.visible = False
             else:
                 q.page["dataset"].progress_bar.visible = False
+                q.page["dataset"].success_bar.text = f"Data successfully uploaded, it has {n_rows:,} rows!"
                 q.page["dataset"].success_bar.visible = True
         except Exception as e:
             logging.error(f"Something went wrong while uploading the dataset: {e}")
@@ -460,7 +462,7 @@ async def datasets(q: Q):
                 ui.message_bar(
                     name="success_bar",
                     type="success",
-                    text=f"Data successfully uploaded, it has {q.args.n_rows} rows!",
+                    text=f"Data successfully uploaded!",
                     visible=False,
                 ),
                 ui.file_upload(
@@ -653,7 +655,7 @@ def upload_demo_examples(q: Q):
         q.user.table_info_path = f"{sample_data_path}/table_info.jsonl"
         q.user.sample_qna_path = None
 
-        n_rows, db_resp = db_setup_api(
+        _, db_resp = db_setup_api(
             db_name=q.user.db_name,
             hostname=q.user.host_name,
             user_name=q.user.user_name,
@@ -665,7 +667,6 @@ def upload_demo_examples(q: Q):
         )
         logging.info(f"DB updated with demo examples: \n {db_resp}")
     q.args.table_dropdown = usr_table_name
-    return n_rows
 
 
 async def on_event(q: Q):
@@ -698,7 +699,7 @@ async def on_event(q: Q):
         logging.info(f"User selected task: {q.args.task_dropdown}")
         q.user.task_dropdown = q.args.task_dropdown
         q.page["task_choice"].task_dropdown.value = q.user.task_dropdown
-        q.args.chatbot = f"Task '{TASK_CHOICE[q.user.task_dropdown]}' selected"
+        q.args.chatbot = f"'{TASK_CHOICE[q.user.task_dropdown]}' mode selected"
         # Refresh response is triggered when user selects a table via dropdown
         event_handled = True
     if (
@@ -763,7 +764,7 @@ async def on_event(q: Q):
     elif q.args.demo_mode:
         logging.info(f"Switching to demo mode!")
         # If demo datasets are not present, register them.
-        _ = upload_demo_examples(q)
+        upload_demo_examples(q)
         logging.info(f"Demo dataset selected: {q.user.table_name}")
         await submit_table(q)
         sample_qs = """