h2oai / sql-sidekick

Experiment on QnA tabular data using LLMs and SQL
Apache License 2.0
23 stars 2 forks source link

Pass Db dialect properly #63

Closed pramitchoudhary closed 9 months ago

pramitchoudhary commented 9 months ago
from sidekick.prompter import db_setup, ask
from sidekick.schema_generator import generate_schema
from sidekick.utils import setup_dir, list_db_dialects

list_db_dialects()

res = ask(
        question=question,
        table_info_path=...,
        sample_queries_path=...,
        table_name=...,
        db_dialect="databricks",
        is_command=False,
        model_name="h2ogpt-sql-nsql-llama-2-7B",
        is_regenerate=...,
        is_regen_with_options=...,
        execute_query=False,
        local_base_path=...
    )

With this change, one can also pass the table schema directly, e.g.

input_schema = [{'Column Name': 'Person_ID', 'Column Type': 'NUMERIC'}, 
{'Column Name': 'Gender', 'Column Type': 'TEXT', 'Sample Values': ['Male', 'Female']},
...
]

db_setup(
          db_name=DB_NAME,
          hostname=HOST_NAME,
          user_name=USER_NAME,
          password=PASSWORD,
          port=PORT,
          table_name=table_name,
          table_schema = input_schema,
          local_base_path=base_path,
          add_sample=False
    )