Refactoring

app.py

@@ -7,26 +7,26 @@ import pandas as pd
 from langchain_huggingface.llms import HuggingFacePipeline
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
 from langsmith import traceable
-
+from langchain import hub
 
 # Height of the Tabs Text Area
 TAB_LINES = 8
-# Load Token
-md_token = os.getenv('MD_TOKEN')
 
-print('Connecting to DB...')
-# Connect to DB
+
+#----------CONNECT TO DATABASE----------
+md_token = os.getenv('MD_TOKEN')
 conn = duckdb.connect(f"md:my_db?motherduck_token={md_token}", read_only=True)
 
+#---------------------------------------
 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 HUGGINGFACE PIPELINE-------
 tokenizer = AutoTokenizer.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1")
 
 quantization_config = BitsAndBytesConfig(
@@ -40,9 +40,15 @@ model = AutoModelForCausalLM.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1",
 
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=1024, return_full_text=False)
 hf = HuggingFacePipeline(pipeline=pipe)
-print('Model Loaded...')
-print(f'Model Device: {model.device}')
+#---------------------------------------
 
+#-----LOAD PROMPT FROM LANCHAIN HUB-----
+prompt = hub.pull("sql-agent-prompt")
+#---------------------------------------
+
+
+
+#--------------ALL UTILS----------------
 # Get Databases
 def get_schemas():
     schemas = conn.execute("""
@@ -78,24 +84,15 @@ def get_table_schema(table):
 
 # Get Prompt
 def get_prompt(schema, query_input):
-    text = f"""
-    ### Instruction:
-    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}
-    
-    ### Question:
-    {query_input}
-    ### Response (use duckdb shorthand if possible):
-    """
-    return text
+    return prompt.format(schema=schema, query_input=query_input)
 
 @spaces.GPU(duration=60)
 @traceable()
 def generate_sql(prompt):
     result = hf.invoke(prompt)
     return result.strip()
+#---------------------------------------
+
 
 # Generate SQL
 def text2sql(table, query_input):

requirements.txt

@@ -3,4 +3,5 @@ bitsandbytes==0.44.1
 transformers==4.44.2
 duckdb==1.1.1
 langsmith==0.1.135
+langchain==0.3.4
 langchain-huggingface