Shankarm08/Text2SQL

1import os2import streamlit as st3from dotenv import load_dotenv4from langchain import HuggingFaceHub5 6# Load environment variables from the .env file7load_dotenv()8 9# Set your Hugging Face API token from the environment variable10HUGGINGFACE_API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")11 12# Function to return the SQL query from natural language input13def load_sql_query(question):14    try:15        # Initialize the Hugging Face model using LangChain's HuggingFaceHub class16        llm = HuggingFaceHub(17            repo_id="Salesforce/grappa_large_jnt",  # Hugging Face model repo for text-to-SQL18            task="text2text-generation",  # Set the task to 'text2text-generation'19            huggingfacehub_api_token=HUGGINGFACE_API_TOKEN,  # Pass your API token20            model_kwargs={"temperature": 0.3}  # Optional: Adjust response randomness21        )22        23        # Call the model with the user's question and get the SQL query24        sql_query = llm.predict(question)25        return sql_query26    except Exception as e:27        # Capture and return any exceptions or errors28        return f"Error: {str(e)}"29 30# Streamlit App UI starts here31st.set_page_config(page_title="Text-to-SQL Demo", page_icon=":robot:")32st.header("Text-to-SQL Demo")33 34# Function to get user input35def get_text():36    input_text = st.text_input("Ask a question (related to a database):", key="input")37    return input_text38 39# Get user input40user_input = get_text()41 42# Create a button for generating the SQL query43submit = st.button('Generate SQL')44 45# If the generate button is clicked and user input is not empty46if submit and user_input:47    response = load_sql_query(user_input)48    st.subheader("Generated SQL Query:")49    st.write(response)50elif submit:51    st.warning("Please enter a question.")  # Warning for empty input52 53