Spaces:
Sleeping
Sleeping
File size: 4,177 Bytes
4851168 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
import gradio as gr
from PIL import Image
import os
import sqlite3
import google.generativeai as genai
import time
# Initialize Gemini
genai.configure(api_key="AIzaSyCnuJMrTkckScv0YxAaaQzH1LDomNUUppA")
genai_model = genai.GenerativeModel('gemini-pro')
class SQLPromptModel:
def __init__(self, database):
self.database = database
self.conn = sqlite3.connect(self.database)
def fetch_table_schema(self, table_name):
cursor = self.conn.cursor()
cursor.execute(f"PRAGMA table_info({table_name})")
schema = cursor.fetchall()
return schema if schema else None
def text2sql_gemini(self, schema, user_prompt, inp_prompt=None):
table_columns = ', '.join([f"{col[1]} {col[2]}" for col in schema])
prompt = f"""Below are SQL table schemas paired with instructions that describe a task.
Using valid SQLite, write a response that appropriately completes the request for the provided tables.
### Instruction: {user_prompt} ###
Input: CREATE TABLE sql_pdf({table_columns});
### Response: (Return only generated query based on user_prompt , nothing extra)"""
if inp_prompt is not None:
prompt = prompt.replace(user_prompt, inp_prompt + " ")
completion = genai_model.generate_content(prompt)
generated_query = completion.text
# Extract SQL query
start_index = generated_query.find("SELECT")
end_index = generated_query.find(";", start_index) + 1
if start_index != -1 and end_index != 0:
return generated_query[start_index:end_index]
return generated_query
def execute_query(self, query):
cur = self.conn.cursor()
cur.execute(query)
columns = [header[0] for header in cur.description]
rows = [row for row in cur.fetchall()]
cur.close()
self.conn.commit()
return rows, columns
def execute_sql_query(input_prompt):
database = r"sql_pdf.db"
sql_model = SQLPromptModel(database)
user_prompt = "Give complete details of properties in India"
for _ in range(3): # Retry logic
try:
table_schema = sql_model.fetch_table_schema("sql_pdf")
if table_schema:
if input_prompt.strip():
query = sql_model.text2sql_gemini(table_schema, user_prompt, input_prompt)
else:
query = sql_model.text2sql_gemini(table_schema, user_prompt, user_prompt)
rows, columns = sql_model.execute_query(query)
return {"Query": query, "Results": rows, "Columns": columns}
else:
return {"error": "Table schema not found."}
except Exception as e:
print(f"An error occurred: {e}")
time.sleep(1)
return {"error": "Failed to execute query after 3 retries."}
# Load the image
image = Image.open(os.path.join(os.path.abspath(''), "house_excel_sheet.png"))
# Create Gradio interface
with gr.Blocks(title="House Database Query") as demo:
gr.Markdown("# House Database Query System")
# Display the image
gr.Image(image)
gr.Markdown("""### The database contains information about different properties including their fundamental details.
You can query this database using natural language.""")
# Query input and output
query_input = gr.Textbox(
lines=2,
label="Database Query",
placeholder="Enter your query or choose from examples below. Default: 'Properties in India'"
)
query_output = gr.JSON(label="Query Results")
# Example queries
gr.Examples(
examples=[
"Properties in France",
"Properties greater than an acre",
"Properties with more than 400 bedrooms"
],
inputs=query_input,
outputs=query_output,
fn=execute_sql_query
)
# Submit button
query_input.submit(fn=execute_sql_query, inputs=query_input, outputs=query_output)
if __name__ == "__main__":
demo.launch(share=True) |