Update app.py

app.py

@@ -1,16 +1,44 @@
-import gradio as gr
-from datasets import load_dataset
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
-# Load datasets
-ds1 = load_dataset("b-mc2/sql-create-context")
-ds2 = load_dataset("TuneIt/o1-python")
-ds3 = load_dataset("HuggingFaceFW/fineweb-2", "aai_Latn")
-ds4 = load_dataset("HuggingFaceFW/fineweb-2", "aai_Latn_removed")
-ds5 = load_dataset("HuggingFaceFW/fineweb-2", "aak_Latn")
-ds6 = load_dataset("sentence-transformers/embedding-training-data")
+# Load the tokenizer and models
+tokenizer = AutoTokenizer.from_pretrained("cssupport/t5-small-awesome-text-to-sql")
+original_model = AutoModelForSeq2SeqLM.from_pretrained("cssupport/t5-small-awesome-text-to-sql", torch_dtype=torch.bfloat16)
+ft_model = AutoModelForSeq2SeqLM.from_pretrained("daljeetsingh/sql_ft_t5small_kag", torch_dtype=torch.bfloat16)
 
-# Load the model and create a Gradio interface
-demo = gr.load("huggingface/distilbert-base-uncased")
+# Move models to GPU
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+original_model.to(device)
+ft_model.to(device)
 
-# Launch the Gradio app
-demo.launch()
+# Streamlit app layout
+st.title("SQL Generation with T5 Models")
+
+# Input text box
+input_text = st.text_area("Enter your query:", height=150)
+
+# Generate button
+if st.button("Generate SQL"):
+    if input_text:
+        # Tokenize input
+        inputs = tokenizer(input_text, return_tensors='pt').to(device)
+
+        # Generate SQL queries
+        with torch.no_grad():
+            original_sql = tokenizer.decode(
+                original_model.generate(inputs["input_ids"], max_new_tokens=200)[0],
+                skip_special_tokens=True
+            )
+            ft_sql = tokenizer.decode(
+                ft_model.generate(inputs["input_ids"], max_new_tokens=200)[0],
+                skip_special_tokens=True
+            )
+
+        # Display results
+        st.subheader("Original Model Output")
+        st.write(original_sql)
+        st.subheader("Fine-Tuned Model Output")
+        st.write(ft_sql)
+    else:
+        st.warning("Please enter a query to generate SQL.")