Spaces:
Sleeping
Sleeping
File size: 2,786 Bytes
37f526a 63e6a46 d29a670 e6e05c7 d29a670 e6e05c7 d29a670 e6e05c7 d29a670 63e6a46 d29a670 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import torch
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.\nReturn only SQL (no text):\nQuestion: {question}\nSQL:"
# ✅ Use .invoke() instead of calling the object directly
response = sqlcoder_llm.invoke(prompt)
# Ensure we get plain string
if isinstance(response, dict) and "text" in response:
response = response["text"]
elif isinstance(response, list):
response = response[0]["generated_text"]
# Clean and finalize SQL
sql_query = response.strip().split("SQL:")[-1].strip()
sql_query = sql_query.split("\n")[0].strip()
if not sql_query.endswith(";"):
sql_query += ";"
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}\n\nGenerated SQL:\n{sql_query}", 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()
|