File size: 1,549 Bytes
0dc11a4 6939e2c 0dc11a4 2758cd1 0dc11a4 2758cd1 0dc11a4 |
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 |
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
# 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)
# Move models to GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
original_model.to(device)
ft_model.to(device)
# 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.") |