bol20162021
commited on
Commit
•
8d201a1
1
Parent(s):
642b13b
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""Streamlit app for demoing nsql-llama-2-70B."""
|
2 |
+
|
3 |
+
import json
|
4 |
+
import os
|
5 |
+
|
6 |
+
import pandas as pd
|
7 |
+
import requests
|
8 |
+
import streamlit as st
|
9 |
+
from manifest import Manifest, Response
|
10 |
+
from manifest.connections.client_pool import ClientConnection
|
11 |
+
|
12 |
+
STOP_TOKENS = ["###", ";", "--", "```"]
|
13 |
+
|
14 |
+
|
15 |
+
def generate_prompt(question, schema):
|
16 |
+
return f"""{schema}\n\n\n-- Using valid SQLite, answer the following questions for the tables provided above.\n\n-- {question}\n"""
|
17 |
+
|
18 |
+
|
19 |
+
def generate_sql(question, schema):
|
20 |
+
prompt = generate_prompt(question, schema)
|
21 |
+
url = st.secrets["backend_url"]
|
22 |
+
headers = {
|
23 |
+
"Content-Type": "application/json",
|
24 |
+
"key": st.secrets["key"],
|
25 |
+
}
|
26 |
+
|
27 |
+
data = {
|
28 |
+
"inputs": [prompt],
|
29 |
+
"params": {
|
30 |
+
"do_sample": {"type": "bool", "value": "false"},
|
31 |
+
"max_tokens_to_generate": {"type": "int", "value": "1000"},
|
32 |
+
"repetition_penalty": {"type": "float", "value": "1"},
|
33 |
+
"temperature": {"type": "float", "value": "1"},
|
34 |
+
"top_k": {"type": "int", "value": "50"},
|
35 |
+
"top_logprobs": {"type": "int", "value": "0"},
|
36 |
+
"top_p": {"type": "float", "value": "1"},
|
37 |
+
},
|
38 |
+
}
|
39 |
+
|
40 |
+
r = requests.post(url, headers=headers, data=json.dumps(data), stream=True)
|
41 |
+
|
42 |
+
if r.encoding is None:
|
43 |
+
r.encoding = "utf-8"
|
44 |
+
for line in r.iter_lines(decode_unicode=True):
|
45 |
+
if line and line.startswith("data: "):
|
46 |
+
output = json.loads(line[len("data: ") :])
|
47 |
+
token = output.get("stream_token", "")
|
48 |
+
if len(token) > 0:
|
49 |
+
yield token
|
50 |
+
|
51 |
+
|
52 |
+
st.title("nsql-llama-2-70B Demo")
|
53 |
+
|
54 |
+
expander = st.expander("Database Schema")
|
55 |
+
|
56 |
+
# Input field for text prompt
|
57 |
+
# TODO(Bo Li): update this with the new example
|
58 |
+
default_schema = """CREATE TABLE stadium (
|
59 |
+
stadium_id number,
|
60 |
+
location text,
|
61 |
+
name text,
|
62 |
+
capacity number,
|
63 |
+
highest number,
|
64 |
+
lowest number,
|
65 |
+
average number
|
66 |
+
)
|
67 |
+
CREATE TABLE singer (
|
68 |
+
singer_id number,
|
69 |
+
name text,
|
70 |
+
country text,
|
71 |
+
song_name text,
|
72 |
+
song_release_year text,
|
73 |
+
age number,
|
74 |
+
is_male others
|
75 |
+
)
|
76 |
+
CREATE TABLE concert (
|
77 |
+
concert_id number,
|
78 |
+
concert_name text,
|
79 |
+
theme text,
|
80 |
+
stadium_id text,
|
81 |
+
year text
|
82 |
+
)
|
83 |
+
CREATE TABLE singer_in_concert (
|
84 |
+
concert_id number,
|
85 |
+
singer_id text
|
86 |
+
)"""
|
87 |
+
|
88 |
+
schema = expander.text_area("Current schema:", value=default_schema, height=500)
|
89 |
+
|
90 |
+
# Input field for text prompt
|
91 |
+
text_prompt = st.text_input(
|
92 |
+
"Please let me know what question do you want to ask?",
|
93 |
+
value="What is the maximum, the average, and the minimum capacity of stadiums ?",
|
94 |
+
)
|
95 |
+
|
96 |
+
|
97 |
+
# if text_prompt or
|
98 |
+
if st.button("Generate SQL"):
|
99 |
+
sql_query = generate_sql(text_prompt, schema)
|
100 |
+
st.write_stream(sql_query)
|