Chemically-motivated/SQL_Generation

1import streamlit as st2import torch3from transformers import AutoModelForSeq2SeqLM, AutoTokenizer4from transformers.utils import logging5 6# Set up logging7logging.set_verbosity_info()8logger = logging.get_logger("transformers")9 10# Model names11original_model_name = 't5-small'12fine_tuned_model_name = 'daljeetsingh/sql_ft_t5small_kag'13 14# Load models and tokenizer15tokenizer = AutoTokenizer.from_pretrained(original_model_name)16original_model = AutoModelForSeq2SeqLM.from_pretrained(original_model_name, torch_dtype=torch.bfloat16)17fine_tuned_model = AutoModelForSeq2SeqLM.from_pretrained(fine_tuned_model_name, torch_dtype=torch.bfloat16)18 19# Move models to GPU20device = 'cuda' if torch.cuda.is_available() else 'cpu'21original_model.to(device)22fine_tuned_model.to(device)23 24def generate_sql_query(prompt):25    """26    Generate SQL queries using both the original and fine-tuned models.27    """28    inputs = tokenizer(prompt, return_tensors='pt').to(device)29    try:30        # Generate output from the original model31        original_output = original_model.generate(32            inputs["input_ids"], 33            max_new_tokens=200,34        )35        original_sql = tokenizer.decode(36            original_output[0], 37            skip_special_tokens=True38        )39 40        # Generate output from the fine-tuned model41        fine_tuned_output = fine_tuned_model.generate(42            inputs["input_ids"], 43            max_new_tokens=200,44        )45        fine_tuned_sql = tokenizer.decode(46            fine_tuned_output[0], 47            skip_special_tokens=True48        )49 50        return original_sql, fine_tuned_sql51    except Exception as e:52        logger.error(f"Error: {str(e)}")53        return f"Error: {str(e)}", None54 55# Streamlit App Interface56st.title("SQL Query Generation")57st.markdown("This application generates SQL queries based on your input prompt.")58 59# Input prompt60prompt = st.text_area(61    "Enter your prompt here...",62    value="Find all employees who joined after 2020.",63    height=15064)65 66# Generate button67if st.button("Generate"):68    if prompt:69        original_sql, fine_tuned_sql = generate_sql_query(prompt)70        st.subheader("Original Model Output")71        st.text_area("Original SQL Query", value=original_sql, height=200)72        st.subheader("Fine-Tuned Model Output")73        st.text_area("Fine-Tuned SQL Query", value=fine_tuned_sql, height=200)74    else:75        st.warning("Please enter a prompt to generate SQL queries.")76 77# Examples78st.sidebar.title("Examples")79st.sidebar.markdown("""80- **Example 1**: Find all employees who joined after 2020.81- **Example 2**: Retrieve the names of customers who purchased product X in the last month.82""")83