Update app.py

app.py

@@ -1,72 +1,67 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
 import streamlit as st
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-@st.cache_resource
-def load_models():
-    question_model_name = "mrm8488/t5-base-finetuned-question-generation-ap"
-    recipe_model_name = "flax-community/t5-recipe-generation"
-    instruct_model_name = "norallm/normistral-7b-warm-instruct"
+# 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 T5-based models for question generation and recipe generation
-    question_model = AutoModelForSeq2SeqLM.from_pretrained(question_model_name)
-    question_tokenizer = AutoTokenizer.from_pretrained(question_model_name)
+sql_tokenizer = AutoTokenizer.from_pretrained("juierror/flan-t5-text2sql-with-schema-v2")
+sql_model = AutoModelForSeq2SeqLM.from_pretrained("juierror/flan-t5-text2sql-with-schema-v2")
 
-    recipe_model = AutoModelForSeq2SeqLM.from_pretrained(recipe_model_name)
-    recipe_tokenizer = AutoTokenizer.from_pretrained(recipe_model_name)
+# 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}"
 
-    # Load the instruction model as a causal language model
-    instruct_model = AutoModelForCausalLM.from_pretrained(instruct_model_name)
-    instruct_tokenizer = AutoTokenizer.from_pretrained(instruct_model_name)
+    # 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
 
-    return (question_model, question_tokenizer), (recipe_model, recipe_tokenizer), (instruct_model, instruct_tokenizer)
+# 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 generate questions using the question model
-def generate_question(text, model, tokenizer):
-    input_text = f"generate question: {text}"
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
-    outputs = model.generate(input_ids)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    input_ids = sql_tokenizer(prompt, max_length=512, return_tensors="pt").input_ids
+    return input_ids
 
-# Function to generate recipes using the recipe model
-def generate_recipe(ingredients, model, tokenizer):
-    input_text = f"generate recipe: {ingredients}"
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
-    outputs = model.generate(input_ids)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+# 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
 
-# Function to generate instructions using the instruction model
-def generate_instruction(prompt, model, tokenizer):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    outputs = model.generate(input_ids)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Streamlit Interface
+# Streamlit UI
 def main():
-    st.title("Multi-Model Application")
-
-    # Load all models
-    (question_model, question_tokenizer), (recipe_model, recipe_tokenizer), (instruct_model, instruct_tokenizer) = load_models()
+    st.title("Multi-Model: Text to SQL and Question Generation")
 
-    # Tabs for different functionalities
-    tab = st.selectbox("Choose task", ["Question Generation", "Recipe Generation", "Instruction Following"])
+    # 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 = {}
 
-    if tab == "Question Generation":
-        passage = st.text_area("Enter a passage for question generation")
-        if st.button("Generate Question"):
-            question = generate_question(passage, question_model, question_tokenizer)
-            st.write("Generated Question:", question)
+    # If tables are provided, generate a question
+    if tables:
+        generated_question = generate_question(tables)
+        st.write(f"Generated Question: {generated_question}")
 
-    elif tab == "Recipe Generation":
-        ingredients = st.text_area("Enter ingredients for recipe generation")
-        if st.button("Generate Recipe"):
-            recipe = generate_recipe(ingredients, recipe_model, recipe_tokenizer)
-            st.write("Generated Recipe:", recipe)
+    # Input question manually if needed
+    question = st.text_area("Enter your question (optional):", generated_question if tables else "")
 
-    elif tab == "Instruction Following":
-        instruction_prompt = st.text_area("Enter an instruction prompt")
-        if st.button("Generate Instruction"):
-            instruction = generate_instruction(instruction_prompt, instruct_model, instruct_tokenizer)
-            st.write("Generated Instruction:", instruction)
+    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.")
 
 if __name__ == "__main__":
     main()