Spaces:
Sleeping
Sleeping
File size: 2,231 Bytes
63e6a46 |
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 |
import sqlite3
import pandas as pd
import gradio as gr
from langchain_community.llms import HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
# ============================================================
# 🚀 Load SQLCoder model
# ============================================================
model_id = "defog/sqlcoder-7b-2"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype="auto",
device_map="auto"
)
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=256,
do_sample=False
)
sqlcoder_llm = HuggingFacePipeline(pipeline=pipe)
# ============================================================
# 🧠 Define query function
# ============================================================
def ask_question(user_db, question):
"""Takes an uploaded SQLite database + a question, returns SQL + result"""
if not user_db:
return "❌ Please upload a database file.", None
conn = sqlite3.connect(user_db.name)
cursor = conn.cursor()
# Create a Text-to-SQL prompt
prompt = f"Translate this question into an SQLite query:\nQuestion: {question}\nSQL:"
sql_query = sqlcoder_llm(prompt)
try:
cursor.execute(sql_query)
rows = cursor.fetchall()
columns = [desc[0] for desc in cursor.description]
df = pd.DataFrame(rows, columns=columns)
conn.close()
return sql_query, df
except Exception as e:
conn.close()
return f"❌ Error executing query: {e}", None
# ============================================================
# 🎨 Gradio UI
# ============================================================
demo = gr.Interface(
fn=ask_question,
inputs=[
gr.File(label="Upload SQLite Database (.db)"),
gr.Textbox(label="Ask your question")
],
outputs=[
gr.Textbox(label="Generated SQL Query"),
gr.Dataframe(label="Query Result")
],
title="🧠 Text-to-SQL on Your Own Database",
description="Upload your SQLite database and ask natural language questions."
)
if __name__ == "__main__":
demo.launch()
|