import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load the models
tokenizer_sql = AutoTokenizer.from_pretrained("juierror/flan-t5-text2sql-with-schema-v2")
model_sql = AutoModelForSeq2SeqLM.from_pretrained("juierror/flan-t5-text2sql-with-schema-v2")

tokenizer_question = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-question-generation-ap")
model_question = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-question-generation-ap")

# Function to create the prompt for SQL model
def get_prompt_sql(tables, question):
    return f"""convert question and table into SQL query. tables: {tables}. question: {question}"""

# Function to prepare input data for the SQL model
def prepare_input_sql(question: str, tables: dict):
    tables = [f"""{table_name}({','.join(tables[table_name])})""" for table_name in tables]
    tables = ", ".join(tables)
    prompt = get_prompt_sql(tables, question)
    input_ids = tokenizer_sql(prompt, max_length=512, return_tensors="pt").input_ids
    return input_ids

# Inference function for the SQL model
def inference_sql(question: str, tables: dict) -> str:
    input_data = prepare_input_sql(question=question, tables=tables)
    input_data = input_data.to(model_sql.device)
    outputs = model_sql.generate(inputs=input_data, num_beams=10, top_k=10, max_length=512)
    return tokenizer_sql.decode(outputs[0], skip_special_tokens=True)

# Function to create the prompt for Question Generation model
def get_prompt_question(context):
    return f"generate a question from the following context: {context}"

# Function to prepare input data for the Question Generation model
def prepare_input_question(context: str):
    prompt = get_prompt_question(context)
    input_ids = tokenizer_question(prompt, max_length=512, return_tensors="pt").input_ids
    return input_ids

# Inference function for the Question Generation model
def inference_question(context: str) -> str:
    input_data = prepare_input_question(context)
    input_data = input_data.to(model_question.device)
    outputs = model_question.generate(inputs=input_data, num_beams=10, top_k=10, max_length=512)
    return tokenizer_question.decode(outputs[0], skip_special_tokens=True)

# Streamlit UI
def main():
    st.title("Multi-Model: Text to SQL and Question Generation")

    # Model selection
    model_choice = st.selectbox("Select a model", ["Text to SQL", "Question Generation"])

    # Input question and table schema for SQL model
    if model_choice == "Text to SQL":
        st.subheader("Text to SQL Model")
        question = st.text_area("Enter your question:")
        tables_input = st.text_area("Enter table schemas (in JSON format):", '{"people_name": ["id", "name"], "people_age": ["people_id", "age"]}')
        try:
            tables = eval(tables_input)  # Convert string to dict safely
        except:
            tables = {}

        if st.button("Generate SQL Query"):
            if question and tables:
                sql_query = inference_sql(question, tables)
                st.write(f"Generated SQL Query: {sql_query}")
            else:
                st.write("Please enter both a question and table schemas.")

    # Input context for Question Generation model
    elif model_choice == "Question Generation":
        st.subheader("Question Generation Model")
        context = st.text_area("Enter context:")
        
        if st.button("Generate Question"):
            if context:
                generated_question = inference_question(context)
                st.write(f"Generated Question: {generated_question}")
            else:
                st.write("Please enter context for question generation.")

if __name__ == "__main__":
    main()