Update app.py

app.py

@@ -1,49 +1,53 @@
+import openai
 import streamlit as st
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
-
-# Load the pre-trained model and tokenizer
-@st.cache_resource
-def load_model():
-    model = GPT2LMHeadModel.from_pretrained("rakeshkiriyath/gpt2Medium_text_to_sql")
-    tokenizer = GPT2Tokenizer.from_pretrained("rakeshkiriyath/gpt2Medium_text_to_sql")
-    return model, tokenizer
-
-# Function to generate SQL from English text
-def generate_text_to_sql(query, model, tokenizer, max_length=256):
-    prompt = f"Translate the following English question to SQL: {query}"
-    
-    # Tokenize the input text
-    input_tensor = tokenizer.encode(prompt, return_tensors='pt').to('cpu')  # Use CPU if no GPU
 
-    # Generate SQL using the model
-    output = model.generate(input_tensor, max_length=max_length, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
+# Set up OpenAI API key
+openai.api_key = "YOUR_API_KEY"  # Replace with your OpenAI API key
+
+# Function to generate SQL query using GPT-3
+def generate_sql_query(query, schema):
+    prompt = f"""
+    Given the following database schema:
+
+    {schema}
+
+    Translate the following question into an SQL query:
+
+    Question: {query}
 
-    # Decode the output
-    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
+    SQL Query:
+    """
+
+    # Request GPT-3 to generate the SQL query
+    response = openai.Completion.create(
+        engine="text-davinci-003",  # Use gpt-3.5-turbo for faster and cheaper results
+        prompt=prompt,
+        max_tokens=150,
+        temperature=0.0,  # Lower temperature for deterministic results
+    )
     
     # Extract and return the SQL query
-    sql_output = decoded_output[len(prompt):].strip()
-    return sql_output
+    sql_query = response.choices[0].text.strip()
+    return sql_query
 
-# Main Streamlit app
+# Streamlit interface
 def main():
     # Display title and instructions
-    st.title("Text-to-SQL with GPT-2")
-    st.write("Enter an English query, and get the corresponding SQL query.")
-
-    # Input box for English query
-    query = st.text_input("Enter your question:")
-
-    if query:
-        # Load the model and tokenizer
-        model, tokenizer = load_model()
-        
-        # Generate SQL query
-        sql_query = generate_text_to_sql(query, model, tokenizer)
-        
-        # Display the SQL query
+    st.title("Text-to-SQL with GPT-3")
+    st.write("Enter a database schema and a natural language query to generate the corresponding SQL query.")
+
+    # Input boxes for schema and query
+    schema_input = st.text_area("Enter the Database Schema:", height=200)
+    query_input = st.text_input("Enter your Question:")
+
+    if schema_input and query_input:
+        # Generate the SQL query
+        sql_query = generate_sql_query(query_input, schema_input)
+
+        # Display the generated SQL query
         st.write("Generated SQL Query:")
         st.code(sql_query)
 
 if __name__ == "__main__":
     main()
+