Using Transformers

app.py

@@ -1,53 +1,43 @@
 import os
+import torch
 import duckdb
 import spaces
 import gradio as gr
 import pandas as pd
-from llama_cpp import Llama
-# from dotenv import load_dotenv
-from huggingface_hub import hf_hub_download
-# load_dotenv()
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+
+
+
 # Height of the Tabs Text Area
 TAB_LINES = 8
 # Load Token
 md_token = os.getenv('MD_TOKEN')
+
+print('Connecting to DB...')
 # Connect to DB
 conn = duckdb.connect(f"md:my_db?motherduck_token={md_token}")
 
-# Custom CSS styling
-custom_css = """
-.gradio-container {
-    background-color: #f0f4f8;
-}
-.logo {
-    max-width: 200px;
-    margin: 20px auto;
-    display: block;
-}
-.gr-button {
-    background-color: #4a90e2 !important;
-}
-.gr-button:hover {
-    background-color: #3a7bc8 !important;
-}
-"""
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    print(f"Using GPU: {torch.cuda.get_device_name(device)}")
+else:
+    device = torch.device("cpu")
+    print("Using CPU")
+
 print('Loading Model...')
-# Load Model
-# @spaces.GPU
-# def load_model():
-llama = Llama(
-        model_path=hf_hub_download(
-        repo_id="motherduckdb/DuckDB-NSQL-7B-v0.1-GGUF",
-        filename="DuckDB-NSQL-7B-v0.1-q8_0.gguf",
-        local_dir='.'
-        ),
-        n_ctx=2048,
-        n_gpu_layers=0
-        )
-    # return llama
-
-# llama = load_model()
+
+tokenizer = AutoTokenizer.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1")
+
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type= "nf4")
+
+model = AutoModelForCausalLM.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1", quantization_config=quantization_config,
+                                             device_map="auto", torch_dtype=torch.bfloat16)
 print('Model Loaded...')
+print(f'Model Device: {model.device}')
 
 # Get Databases
 def get_databases():
@@ -76,7 +66,7 @@ def get_schema(table):
 def get_prompt(schema, query_input):
     text = f"""
     ### Instruction:
-    Your task is to generate valid duckdb SQL to answer the following question.
+    Your task is to generate valid duckdb SQL query to answer the following question.
     ### Input:
     Here is the database schema that the SQL query will run on:
     {schema}
@@ -88,12 +78,7 @@ def get_prompt(schema, query_input):
     return text
 
 # Generate SQL
-# @spaces.GPU
-def generate_sql(prompt):
-    
-    result = llama(prompt, temperature=0.1, max_tokens=1000)
-    return result["choices"][0]["text"]
-
+@spaces.GPU
 def text2sql(table, query_input):
     if table is None:
         return {
@@ -102,11 +87,18 @@ def text2sql(table, query_input):
             generated_query: "",
             result_output:pd.DataFrame([{"error": f"❌ Unable to get the SQL query based on the text. {e}"}])
         }
+
     schema = get_schema(table)
+    print(f'Schema Generated...')
     prompt = get_prompt(schema, query_input)
-    
+    print(f'Prompt Generated...')
     try:
-        result = generate_sql(prompt)
+        print(f'Generating SQL... {model.device}')
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+        input_token_len = input_ids.shape[1]
+        outputs = model.generate(input_ids.to(model.device), max_new_tokens=1024)
+        result = tokenizer.decode(outputs[0][input_token_len:], skip_special_tokens=True)
+        print('SQL Generated...')
     except Exception as e:
         return {
             table_schema: schema,
@@ -116,7 +108,6 @@ def text2sql(table, query_input):
         }
     try:
         query_result = conn.sql(result).df()
-        conn.close()
         
     except Exception as e:
         return {
@@ -126,7 +117,6 @@ def text2sql(table, query_input):
             result_output:pd.DataFrame([{"error": f"❌ Unable to get the SQL query based on the text. {e}"}])
         }
         
-    conn.close()
     return {
         table_schema: schema,
         input_prompt: prompt,
@@ -137,6 +127,24 @@ def text2sql(table, query_input):
 # Load Databases Names
 databases = get_databases()
 
+# Custom CSS styling
+custom_css = """
+.gradio-container {
+    background-color: #f0f4f8;
+}
+.logo {
+    max-width: 200px;
+    margin: 20px auto;
+    display: block;
+}
+.gr-button {
+    background-color: #4a90e2 !important;
+}
+.gr-button:hover {
+    background-color: #3a7bc8 !important;
+}
+"""
+
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"), css=custom_css) as demo:
     gr.Image("logo.png", label=None, show_label=False, container=False, height=100)
     
@@ -168,8 +176,8 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"
                 with gr.Tab("Schema"):
                     table_schema = gr.Textbox(lines=TAB_LINES, label="Schema", value="", interactive=False)
  
-    database_dropdown.change(update_tables, inputs=database_dropdown, outputs=tables_dropdown)
-    generate_query_button.click(text2sql, inputs=[tables_dropdown, query_input], outputs=[table_schema, input_prompt, generated_query, result_output])
+        database_dropdown.change(update_tables, inputs=database_dropdown, outputs=tables_dropdown)
+        generate_query_button.click(text2sql, inputs=[tables_dropdown, query_input], outputs=[table_schema, input_prompt, generated_query, result_output])
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -5,5 +5,7 @@ huggingface_hub
 python-dotenv
 scikit-build-core
 duckdb
-https://github.com/abetlen/llama-cpp-python/releases/download/v0.2.82-cu124/llama_cpp_python-0.2.82-cp310-cp310-linux_x86_64.whl
 gradio
+transformers
+bitsandbytes
+torch