bol20162021 commited on
Commit
8d201a1
1 Parent(s): 642b13b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +100 -0
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)