Enable task mode #44

Author: pramitchoudhary

sidekick/prompter.py

@@ -452,6 +452,9 @@ def query_api(
             logger.info("Executing user provided SQL without re-generation...")
             res = question.strip().lower().split("execute sql:")[1].strip()
         else:
+            _check_cond = question.strip().lower().split("execute sql:")
+            if len(_check_cond) > 1:
+                question = question.strip().lower().split("execute sql:")[1].strip()
             res, alt_res = sql_g.generate_sql(table_names, question, model_name=model_name, _dialect=db_dialect)
         logger.info(f"Input query: {question}")
         logger.info(f"Generated response:\n\n{res}")

sidekick/utils.py

@@ -15,18 +15,23 @@
 from sentence_transformers import SentenceTransformer
 from sidekick.logger import logger
 from sklearn.metrics.pairwise import cosine_similarity
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import (AutoConfig, AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig)
 
-model_choices_map = {
+MODEL_CHOICE_MAP = {
     "h2ogpt-sql-sqlcoder2": "defog/sqlcoder2",
     "h2ogpt-sql-nsql-llama-2-7B": "NumbersStation/nsql-llama-2-7B",
 }
 
-model_device_map = {
+MODEL_DEVICE_MAP = {
     "h2ogpt-sql-sqlcoder2": 0,
     "h2ogpt-sql-nsql-llama-2-7B": 1,
 }
 
+TASK_CHOICE = {
+    "q_a": "Question/Answering",
+    "sqld": "SQL Debugging",
+}
 
 def generate_sentence_embeddings(model_path: str, x, batch_size: int = 32, device: Optional[str] = None):
     # Reference:
@@ -290,7 +295,7 @@ def is_resource_low(model_name: str):
         off_load = False
     else:
         n_gpus = torch.cuda.device_count()
-        device_index = model_device_map[model_name] if model_name and n_gpus > 1 else 0
+        device_index = MODEL_DEVICE_MAP[model_name] if model_name and n_gpus > 1 else 0
         logger.debug(f"Information on device: {device_index}")
         free_in_GB = int(torch.cuda.mem_get_info(device_index)[0] / 1024**3)
         total_memory = int(torch.cuda.get_device_properties(device_index).total_memory / 1024**3)
@@ -382,14 +387,14 @@ def _load_llm(model_type: str, device_index: int = 0, load_in_4bit=True):
 
         if not model_type:  # if None, load all models
             for device_index in range(n_gpus):
-                model_name = list(model_choices_map.values())[device_index]
+                model_name = list(MODEL_CHOICE_MAP.values())[device_index]
                 model, tokenizer = _load_llm(model_name, device_index)
-                _name = list(model_choices_map.keys())[device_index]
+                _name = list(MODEL_CHOICE_MAP.keys())[device_index]
                 models[_name] = model
                 tokenizers[_name] = tokenizer
         else:
-            model_name = model_choices_map[model_type]
-            d_index = model_device_map[model_type] if n_gpus > 1 else 0
+            model_name = MODEL_CHOICE_MAP[model_type]
+            d_index = MODEL_DEVICE_MAP[model_type] if n_gpus > 1 else 0
             model, tokenizer = _load_llm(model_name, d_index)
             models[model_type] = model
             tokenizers[model_type] = tokenizer

ui/app.py

@@ -12,7 +12,7 @@
 from h2o_wave.core import expando_to_dict
 from sidekick.prompter import db_setup_api, query_api
 from sidekick.query import SQLGenerator
-from sidekick.utils import get_table_keys, save_query, setup_dir, update_tables
+from sidekick.utils import get_table_keys, save_query, setup_dir, update_tables, TASK_CHOICE
 
 # Load the config file and initialize required paths
 base_path = (Path(__file__).parent / "../").resolve()
@@ -91,7 +91,7 @@ def clear_cards(q, ignore: Optional[List[str]] = []) -> None:
 async def chat(q: Q):
     q.page["sidebar"].value = "#chat"
 
-    if q.args.table_dropdown or q.args.model_choice_dropdown:
+    if q.args.table_dropdown or q.args.model_choice_dropdown or q.args.task_dropdown:
         # If a table/model is selected, the trigger causes refresh of the page
         # so we update chat history with table name selection and return
         # avoiding re-drawing.
@@ -113,6 +113,9 @@ async def chat(q: Q):
         ui.choice("h2ogpt-sql-sqlcoder2", "h2ogpt-sql-sqlcoder2"),
     ]
     q.user.model_choice_dropdown = "h2ogpt-sql-sqlcoder2"
+
+    task_choices = [ui.choice("q_a", "Question/Answering"), ui.choice("sqld", "SQL Debugging")]
+    q.user.task_choice_dropdown = "q_a"
     add_card(
         q,
         "background_card",
@@ -123,7 +126,7 @@ async def chat(q: Q):
                 ui.inline(items=[ui.toggle(name="demo_mode", label="Demo", trigger=True)], justify="end"),
             ],
         ),
-    )
+    ),
 
     add_card(
         q,
@@ -149,7 +152,24 @@ async def chat(q: Q):
                 ),
             ],
         ),
-    )
+    ),
+    add_card(
+        q,
+        "task_choice",
+        ui.form_card(
+            box="vertical",
+            items=[
+                ui.dropdown(
+                    name="task_dropdown",
+                    label="Task",
+                    required=True,
+                    choices=task_choices,
+                    value=q.user.task_choice_dropdown if q.user.task_choice_dropdown else None,
+                    trigger=True,
+                )
+            ],
+        ),
+    ),
     add_card(
         q,
         "chat_card",
@@ -228,11 +248,15 @@ 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
     ):
         return
 
     # Append bot response.
     question = f"{q.args.chatbot}"
+    # Check on task choice.
+    if q.user.task_dropdown == "sqld":
+        question = f"Execute SQL:\n{q.args.chatbot}"
     logging.info(f"Question: {question}")
 
     # For regeneration, currently there are 2 modes
@@ -638,7 +662,13 @@ async def on_event(q: Q):
         q.args.chatbot = f"Model {q.user.model_choice_dropdown} selected"
         # Refresh response is triggered when user selects a table via dropdown
         event_handled = True
-
+    if q.args.task_dropdown and not q.args.chatbot and q.user.task_dropdown != q.args.task_dropdown:
+        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"
+        # Refresh response is triggered when user selects a table via dropdown
+        event_handled = True
     if (
         q.args.save_conversation
         or q.args.save_rejected_conversation