update

app.py

@@ -26,7 +26,6 @@ db_schemas = {
         "Citizenship" text,
         PRIMARY KEY ("Singer_ID")
     );
-
     CREATE TABLE "song" (
         "Song_ID" int,
         "Title" text,
@@ -64,4 +63,4 @@ if question:
     # Get the model's response (in this case, SQL query or logits)
     response = text2sql_bot.get_response(question, db_id)
     st.write(f"**Database:** {db_id}")
-    st.write(f"**Results:** {response}")
+    st.write(f"**Results:** {response}")

build_whoosh_index.py

@@ -8,7 +8,6 @@ def extract_contents_from_db(db_path, max_len=25):
     """
     Extract all non-null, unique text values of length <= max_len
     from every table and column in the SQLite database.
-
     Returns:
         List of tuples [(doc_id, text), ...]
     """
@@ -45,7 +44,6 @@ def extract_contents_from_db(db_path, max_len=25):
 def build_index_for_db(db_id, db_path, index_root="db_contents_index"):
     """
     Build (or open) a Whoosh index for a single database.
-
     - If the index already exists in index_root/db_id, it will be opened.
     - Otherwise, a new index is created and populated from the SQLite file.
     """
@@ -89,4 +87,4 @@ if __name__ == "__main__":
             print(f"Building Whoosh index for {db_id}...")
             build_index_for_db(db_id, db_file, INDEX_ROOT)
 
-    print("All indexes built successfully.")
+    print("All indexes built successfully.")

schema_item_filter.py

@@ -347,4 +347,4 @@ if __name__ == "__main__":
     import json
     dataset = json.load(open("./data/sft_eval_{}_text2sql.json".format(dataset_name)))
     
-    sic.evaluate_coverage(dataset)
+    sic.evaluate_coverage(dataset)

text2sql.py

@@ -1,255 +1,180 @@
-# Attribution: Original code by Ruoxin Wang
-# Repository: <your-repo-url>
-
-"""
-Module: refactored_chatbot
-This module provides utilities for loading database schemas, extracting DDL,
-indexing content, and a ChatBot class to generate SQL queries from natural language.
-"""
 import os
 import json
+import torch
+import copy
 import re
+import sqlparse
 import sqlite3
-import copy
-from tqdm import tqdm
 
-import torch
+from tqdm import tqdm
+from utils.db_utils import get_db_schema
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from whoosh import index
-import sqlparse
-
-from utils.db_utils import (
-    get_db_schema,
-    check_sql_executability,
-    get_matched_contents,
-    get_db_schema_sequence,
-    get_matched_content_sequence
-)
+from whoosh.index import create_in
+from whoosh.fields import Schema, TEXT
+from whoosh.qparser import QueryParser
+from utils.db_utils import check_sql_executability, get_matched_contents, get_db_schema_sequence, get_matched_content_sequence
 from schema_item_filter import SchemaItemClassifierInference, filter_schema
 
-
-class DatabaseUtils:
-    """
-    Utilities for loading database comments, schemas, and DDL statements.
-    """
-
-    @staticmethod
-    def _remove_similar_comments(names, comments):
-        """
-        Remove comments identical to table/column names (ignoring underscores and spaces).
-        """
-        filtered = []
-        for name, comment in zip(names, comments):
-            normalized_name = name.replace("_", "").replace(" ", "").lower()
-            normalized_comment = comment.replace("_", "").replace(" ", "").lower()
-            filtered.append("") if normalized_name == normalized_comment else filtered.append(comment)
-        return filtered
-
-    @staticmethod
-    def load_db_comments(table_json_path):
-        """
-        Load additional comments for tables and columns from a JSON file.
-
-        Args:
-            table_json_path (str): Path to JSON file containing table and column comments.
-
-        Returns:
-            dict: Mapping from database ID to comments structure.
-        """
-        additional_info = json.load(open(table_json_path))
-        db_comments = {}
-
-        for db_info in additional_info:
-            db_id = db_info["db_id"]
-            comment_dict = {}
-
-            # Process column comments
-            original_cols = db_info["column_names_original"]
-            col_names = [col.lower() for _, col in original_cols]
-            col_comments = [c.lower() for _, c in db_info["column_names"]]
-            col_comments = DatabaseUtils._remove_similar_comments(col_names, col_comments)
-            col_table_idxs = [t_idx for t_idx, _ in original_cols]
-
-            # Process table comments
-            original_tables = db_info["table_names_original"]
-            tbl_names = [tbl.lower() for tbl in original_tables]
-            tbl_comments = [c.lower() for c in db_info["table_names"]]
-            tbl_comments = DatabaseUtils._remove_similar_comments(tbl_names, tbl_comments)
-
-            for idx, name in enumerate(tbl_names):
-                comment_dict[name] = {
-                    "table_comment": tbl_comments[idx],
-                    "column_comments": {}
-                }
-                # Associate columns
-                for t_idx, col_name, col_comment in zip(col_table_idxs, col_names, col_comments):
-                    if t_idx == idx:
-                        comment_dict[name]["column_comments"][col_name] = col_comment
-
-            db_comments[db_id] = comment_dict
-
-        return db_comments
-
-    @staticmethod
-    def get_db_schemas(db_path, tables_json):
-        """
-        Build a mapping from database ID to its schema representation.
-
-        Args:
-            db_path (str): Directory containing database subdirectories.
-            tables_json (str): Path to JSON with table comments.
-
-        Returns:
-            dict: Mapping from db_id to schema object.
-        """
-        comments = DatabaseUtils.load_db_comments(tables_json)
-        schemas = {}
-        for db_id in tqdm(os.listdir(db_path), desc="Loading schemas"):
-            sqlite_path = os.path.join(db_path, db_id, f"{db_id}.sqlite")
-            schemas[db_id] = get_db_schema(sqlite_path, comments, db_id)
-        return schemas
-
-    @staticmethod
-    def get_db_ddls(db_path):
-        """
-        Extract formatted DDL statements for all tables in each database.
-
-        Args:
-            db_path (str): Directory containing database subdirectories.
-
-        Returns:
-            dict: Mapping from db_id to its DDL string.
-        """
-        ddls = {}
-        for db_id in os.listdir(db_path):
-            conn = sqlite3.connect(os.path.join(db_path, db_id, f"{db_id}.sqlite"))
-            cursor = conn.cursor()
-            cursor.execute("SELECT name, sql FROM sqlite_master WHERE type='table';")
-            ddl_statements = []
-
-            for name, raw_sql in cursor.fetchall():
-                sql = raw_sql or ""
-                sql = re.sub(r'--.*', '', sql).replace("\t", " ")
-                sql = re.sub(r" +", " ", sql)
-                formatted = sqlparse.format(
-                    sql,
-                    keyword_case="upper",
-                    identifier_case="lower",
-                    reindent_aligned=True
-                )
-                # Adjust spacing for readability
-                formatted = formatted.replace(", ", ",\n    ")
-                if formatted.rstrip().endswith(";"):
-                    formatted = formatted.rstrip()[:-1] + "\n);"
-                formatted = re.sub(r"(CREATE TABLE.*?)\(", r"\1(\n    ", formatted)
-                ddl_statements.append(formatted)
-
-            ddls[db_id] = "\n\n".join(ddl_statements)
-        return ddls
-
-
-class ChatBot:
-    """
-    ChatBot for generating and executing SQL queries using a causal language model.
-    """
-
-    def __init__(self, model_name: str = "seeklhy/codes-1b", device: str = "cuda:0") -> None:
-        """
-        Initialize the ChatBot with model and tokenizer.
-
-        Args:
-            model_name (str): HuggingFace model identifier.
-            device (str): CUDA device string or 'cpu'.
-        """
-        os.environ["CUDA_VISIBLE_DEVICES"] = device
+def remove_similar_comments(names, comments):
+    '''
+    Remove table (or column) comments that have a high degree of similarity with their names
+    '''
+    new_comments = []
+    for name, comment in zip(names, comments):    
+        if name.replace("_", "").replace(" ", "") == comment.replace("_", "").replace(" ", ""):
+            new_comments.append("")
+        else:
+            new_comments.append(comment)
+    
+    return new_comments
+
+def load_db_comments(table_json_path):
+    additional_db_info = json.load(open(table_json_path))
+    db_comments = dict()
+    for db_info in additional_db_info:
+        comment_dict = dict()
+
+        column_names = [column_name.lower() for _, column_name in db_info["column_names_original"]]
+        table_idx_of_each_column = [t_idx for t_idx, _ in db_info["column_names_original"]]
+        column_comments = [column_comment.lower() for _, column_comment in db_info["column_names"]]
+        
+        assert len(column_names) == len(column_comments)
+        column_comments = remove_similar_comments(column_names, column_comments)
+
+        table_names = [table_name.lower() for table_name in db_info["table_names_original"]]
+        table_comments = [table_comment.lower() for table_comment in db_info["table_names"]]
+        
+        assert len(table_names) == len(table_comments)
+        table_comments = remove_similar_comments(table_names, table_comments)
+
+        for table_idx, (table_name, table_comment) in enumerate(zip(table_names, table_comments)):
+            comment_dict[table_name] = {
+                "table_comment": table_comment,
+                "column_comments": dict()
+            }
+            for t_idx, column_name, column_comment in zip(table_idx_of_each_column, column_names, column_comments):
+                if t_idx == table_idx:
+                    comment_dict[table_name]["column_comments"][column_name] = column_comment
+
+        db_comments[db_info["db_id"]] = comment_dict
+    
+    return db_comments
+
+def get_db_id2schema(db_path, tables_json):
+    db_comments = load_db_comments(tables_json)
+    db_id2schema = dict()
+
+    for db_id in tqdm(os.listdir(db_path)):
+        db_id2schema[db_id] = get_db_schema(os.path.join(db_path, db_id, db_id + ".sqlite"), db_comments, db_id)
+    
+    return db_id2schema
+
+def get_db_id2ddl(db_path):
+    db_ids = os.listdir(db_path)
+    db_id2ddl = dict()
+
+    for db_id in db_ids:
+        conn = sqlite3.connect(os.path.join(db_path, db_id, db_id + ".sqlite"))
+        cursor = conn.cursor()
+        cursor.execute("SELECT name, sql FROM sqlite_master WHERE type='table';")
+        tables = cursor.fetchall()
+        ddl = []
+        
+        for table in tables:
+            table_name = table[0]
+            table_ddl = table[1]
+            table_ddl.replace("\t", " ")
+            while "  " in table_ddl:
+                table_ddl = table_ddl.replace("  ", " ")
+            
+            table_ddl = re.sub(r'--.*', '', table_ddl)
+            table_ddl = sqlparse.format(table_ddl, keyword_case = "upper", identifier_case = "lower", reindent_aligned = True)
+            table_ddl = table_ddl.replace(", ", ",\n    ")
+            
+            if table_ddl.endswith(";"):
+                table_ddl = table_ddl[:-1]
+            table_ddl = table_ddl[:-1] + "\n);"
+            table_ddl = re.sub(r"(CREATE TABLE.*?)\(", r"\1(\n    ", table_ddl)
+
+            ddl.append(table_ddl)
+        db_id2ddl[db_id] = "\n\n".join(ddl)
+    
+    return db_id2ddl
+
+class ChatBot():
+    def __init__(self) -> None:
+        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+        model_name = "seeklhy/codes-1b"
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            device_map="auto",
-            torch_dtype=torch.float16
-        )
+        self.model = AutoModelForCausalLM.from_pretrained(model_name, device_map = "auto", torch_dtype = torch.float16)
         self.max_length = 4096
         self.max_new_tokens = 256
         self.max_prefix_length = self.max_length - self.max_new_tokens
 
-        # Schema item classifier
-        self.schema_classifier = SchemaItemClassifierInference("Roxanne-WANG/LangSQL")
-
-        # Initialize content searchers
-        self.content_searchers = {}
-        index_dir = "db_contents_index"
-        for db_id in os.listdir(index_dir):
-            path = os.path.join(index_dir, db_id)
-            if index.exists_in(path):
-                self.content_searchers[db_id] = index.open_dir(path).searcher()
+        # Directly loading the model from Hugging Face
+        self.sic = SchemaItemClassifierInference("Roxanne-WANG/LangSQL")
+        self.db_id2content_searcher = dict()
+        for db_id in os.listdir("db_contents_index"):
+            index_dir = os.path.join("db_contents_index", db_id)
+
+            # Open existing Whoosh index directory
+            if index.exists_in(index_dir):
+                ix = index.open_dir(index_dir)
+                # keep a searcher around for querying
+                self.db_id2content_searcher[db_id] = ix.searcher()
             else:
-                raise FileNotFoundError(f"Whoosh index not found for '{db_id}' at '{path}'")
+                raise ValueError(f"No Whoosh index found for '{db_id}' at '{index_dir}'")
 
-        # Load schemas and DDLs
         self.db_ids = sorted(os.listdir("databases"))
-        self.schemas = DatabaseUtils.get_db_schemas("databases", "data/tables.json")
-        self.ddls = DatabaseUtils.get_db_ddls("databases")
-
-    def get_response(self, question: str, db_id: str) -> str:
-        """
-        Generate an executable SQL query for a natural language question.
-
-        Args:
-            question (str): User question in natural language.
-            db_id (str): Identifier of the target database.
-
-        Returns:
-            str: Executable SQL query or an error message.
-        """
-        # Prepare data
-        schema = copy.deepcopy(self.schemas[db_id])
-        contents = get_matched_contents(question, self.content_searchers[db_id])
+        self.db_id2schema = get_db_id2schema("databases", "data/tables.json")
+        self.db_id2ddl = get_db_id2ddl("databases")
+
+    def get_response(self, question, db_id):
         data = {
             "text": question,
-            "schema": schema,
-            "matched_contents": contents
+            "schema": copy.deepcopy(self.db_id2schema[db_id]),
+            "matched_contents": get_matched_contents(question, self.db_id2content_searcher[db_id])
         }
-        data = filter_schema(data, self.schema_classifier, top_k=6, top_m=10)
+        data = filter_schema(data, self.sic, 6, 10)
         data["schema_sequence"] = get_db_schema_sequence(data["schema"])
         data["content_sequence"] = get_matched_content_sequence(data["matched_contents"])
-
-        prefix = (
-            f"{data['schema_sequence']}\n"
-            f"{data['content_sequence']}\n"
-            f"{question}\n"
-        )
-
-        # Tokenize and ensure length limits
-        input_ids = [self.tokenizer.bos_token_id] + self.tokenizer(prefix)["input_ids"]
+        
+        prefix_seq = data["schema_sequence"] + "\n" + data["content_sequence"] + "\n" + data["text"] + "\n"
+        print(prefix_seq)
+        
+        input_ids = [self.tokenizer.bos_token_id] + self.tokenizer(prefix_seq , truncation = False)["input_ids"]
         if len(input_ids) > self.max_prefix_length:
-            input_ids = [self.tokenizer.bos_token_id] + input_ids[-(self.max_prefix_length - 1):]
+            print("the current input sequence exceeds the max_tokens, we will truncate it.")
+            input_ids = [self.tokenizer.bos_token_id] + input_ids[-(self.max_prefix_length-1):]
         attention_mask = [1] * len(input_ids)
-
+        
         inputs = {
-            "input_ids": torch.tensor([input_ids], dtype=torch.int64).to(self.model.device),
-            "attention_mask": torch.tensor([attention_mask], dtype=torch.int64).to(self.model.device)
+            "input_ids": torch.tensor([input_ids], dtype = torch.int64).to(self.model.device),
+            "attention_mask": torch.tensor([attention_mask], dtype = torch.int64).to(self.model.device)
         }
+        input_length = inputs["input_ids"].shape[1]
 
         with torch.no_grad():
-            outputs = self.model.generate(
+            generate_ids = self.model.generate(
                 **inputs,
-                max_new_tokens=self.max_new_tokens,
-                num_beams=4,
-                num_return_sequences=4
+                max_new_tokens = self.max_new_tokens,
+                num_beams = 4,
+                num_return_sequences = 4
             )
 
-        # Decode and choose executable SQL
-        decoded = self.tokenizer.batch_decode(
-            outputs[:, inputs['input_ids'].shape[1]:],
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=False
-        )
-        final_sql = None
-        for sql in decoded:
-            if check_sql_executability(sql, os.path.join("databases", db_id, f"{db_id}.sqlite")) is None:
-                final_sql = sql.strip()
+        generated_sqls = self.tokenizer.batch_decode(generate_ids[:, input_length:], skip_special_tokens = True, clean_up_tokenization_spaces = False)
+        final_generated_sql = None
+        for generated_sql in generated_sqls:
+            execution_error = check_sql_executability(generated_sql, os.path.join("databases", db_id, db_id + ".sqlite"))
+            if execution_error is None:
+                final_generated_sql = generated_sql
                 break
-        if not final_sql:
-            final_sql = decoded[0].strip() or "Sorry, I cannot generate a suitable SQL query."
 
-        return final_sql
+        if final_generated_sql is None:
+            if generated_sqls[0].strip() != "":
+                final_generated_sql = generated_sqls[0].strip()
+            else:
+                final_generated_sql = "Sorry, I can not generate a suitable SQL query for your question."
+        
+        return final_generated_sql.replace("\n", " ")