update

app.py

@@ -1,30 +1,44 @@
 import streamlit as st
 import requests
 
-def generate_prompt(question, schema):
-    instruction = f"""Your task is to generate valid duckdb SQL to answer the following question{"" if (schema == "") else ", given a duckdb database schema."}"""
-    input_text = f"""
-    {schema}
-    
-    Generate a SQL query that answers the question `{question}`.
-    """
-    return f"""### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n### Response:\n"""
 
+PROMPT_TEMPLATE = """### Instruction:\n{instruction}\n\n### Input:\n{input}{context}\n### Question:\n{question}\n\n### Response:\n"""
+INSTRUCTION_TEMPLATE = """Your task is to generate valid duckdb SQL to answer the following question{has_schema}"""  # noqa: E501
+
+def generate_prompt(question, schema):
+    input = ""
+    if schema:
+        input = """Here is the database schema that the SQL query will run on:\n{schema}\n""".format(  # noqa: E501
+            schema=schema
+        )
+    prompt = PROMPT_TEMPLATE.format(
+        instruction = INSTRUCTION_TEMPLATE.format(
+            has_schema="."
+            if schema == ""
+            else ", given a duckdb database schema."
+        ),
+        context="",
+        input=input,
+        question=question + ". Use DuckDB shorthand if possible.",
+    )
+    return prompt
 
 def generate_sql(question, schema):
     prompt = generate_prompt(question, schema)
     s = requests.Session()
-    api_base = "https://text-motherduck-cl7b-chat-fp16-4vycuix6qcp2.octoai.run/v1"
-    url = f"{api_base}/completions"
+    api_base = "https://text-motherduck-sql-fp16-4vycuix6qcp2.octoai.run"
+    url = f"{api_base}/v1/completions"
     body = {
         "model": "motherduck-sql-fp16",
         "prompt": prompt,
         "temperature": 0.1,
         "max_tokens": 200,
-        "stop":'<s>'
+        "stop":'<s>',
+        "n": 1
     }
+    headers = {"Authorization": f"Bearer {st.secrets['octoml_token']}"}
 
-    with s.post(url, json=body) as resp:
+    with s.post(url, json=body, headers=headers) as resp:
         return resp.json()["choices"][0]["text"]
 
 st.title("DuckDB-NSQL-7B Demo")