import streamlit as st
import torch
from transformers import LLMForConditionalGeneration, LLMTokenizer
import sqlite3

# Load Hugging Face LLM2 model and tokenizer
model_name = "microsoft/CodeGPT-small-py"
tokenizer = LLMTokenizer.from_pretrained(model_name)
model = LLMForConditionalGeneration.from_pretrained(model_name)

# Function to generate SQL query
def generate_sql_query(text):
    input_ids = tokenizer.encode(text, return_tensors="pt")
    outputs = model.generate(input_ids, max_length=100, do_sample=False)
    generated_sql = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_sql

# Function to execute SQL query and retrieve results from the database
def execute_query(sql_query):
    conn = sqlite3.connect('C:/Users/Chovatiya.Parth/Desktop/SQL/superstore Creation.sql')
    cursor = conn.cursor()
    cursor.execute(sql_query)
    results = cursor.fetchall()
    conn.close()
    return results

# Streamlit UI
def main():
    st.title("SQL Chatbot")

    user_query = st.text_input("Enter your query:")

    if st.button("Submit"):
        sql_query = generate_sql_query(user_query)
        st.write("Generated SQL query:", sql_query)

        try:
            results = execute_query(sql_query)
            st.write("Results from the database:")
            for row in results:
                st.write(row)
        except Exception as e:
            st.error("An error occurred while executing the SQL query.")
            st.error(e)

if __name__ == "__main__":
    main()