Update app.py

app.py

@@ -1,67 +1,84 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Load the models and tokenizers
-question_tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-question-generation-ap")
-question_model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-question-generation-ap")
+# 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")
 
-sql_tokenizer = AutoTokenizer.from_pretrained("juierror/flan-t5-text2sql-with-schema-v2")
-sql_model = 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 generate a question based on a table schema
-def generate_question(tables):
-    # Convert table schema to string
-    table_str = ", ".join([f"{table}: ({', '.join(cols)})" for table, cols in tables.items()])
-    prompt = f"Generate a question based on the following table schema: {table_str}"
+# 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}"""
 
-    # Tokenize input and generate question
-    input_ids = question_tokenizer(prompt, return_tensors="pt").input_ids
-    output = question_model.generate(input_ids, num_beams=5, max_length=50)
-    question = question_tokenizer.decode(output[0], skip_special_tokens=True)
-    return 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 prepare input data for SQL generation
-def prepare_sql_input(question, tables):
-    table_str = ", ".join([f"{table}({', '.join(cols)})" for table, cols in tables.items()])
-    prompt = f"Convert the question and table schema into an SQL query. Tables: {table_str}. Question: {question}"
+# Function to create the prompt for Question Generation model
+def get_prompt_question(context):
+    return f"generate a question from the following context: {context}"
 
-    input_ids = sql_tokenizer(prompt, max_length=512, return_tensors="pt").input_ids
+# 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 SQL generation
-def generate_sql(question, tables):
-    input_data = prepare_sql_input(question, tables)
-    input_data = input_data.to(sql_model.device)
-    outputs = sql_model.generate(inputs=input_data, num_beams=10, top_k=10, max_length=512)
-    sql_query = sql_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return sql_query
+# 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")
 
-    # Input table schema
-    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 = {}
+    # Model selection
+    model_choice = st.selectbox("Select a model", ["Text to SQL", "Question Generation"])
 
-    # If tables are provided, generate a question
-    if tables:
-        generated_question = generate_question(tables)
-        st.write(f"Generated Question: {generated_question}")
+    # 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 = {}
 
-    # Input question manually if needed
-    question = st.text_area("Enter your question (optional):", generated_question if tables else "")
+        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.")
 
-    if st.button("Generate SQL Query"):
-        if question and tables:
-            sql_query = generate_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()