Spaces:
Build error
Build error
| import streamlit as st | |
| import weaviate_utils | |
| import tapas_utils | |
| def display_initial_buttons(): | |
| if "upload_flow" not in st.session_state: | |
| st.session_state.upload_flow = False | |
| if "query_flow" not in st.session_state: | |
| st.session_state.query_flow = False | |
| if st.button("Upload new CSV"): | |
| st.session_state.upload_flow = True | |
| st.session_state.query_flow = False | |
| if st.button("Query existing data"): | |
| st.session_state.query_flow = True | |
| st.session_state.upload_flow = False | |
| def display_class_dropdown(client): | |
| if st.session_state.upload_flow: | |
| existing_classes = [cls["class"] for cls in client.schema.get()["classes"]] | |
| class_options = existing_classes + ["New Class"] | |
| return st.selectbox("Select a class or create a new one:", class_options, key="class_selector_upload") | |
| elif st.session_state.query_flow: | |
| existing_classes = [cls["class"] for cls in client.schema.get()["classes"]] | |
| class_options = existing_classes + ["Query all data"] | |
| selected_option = st.selectbox("Select a class or query all data:", class_options, key="class_selector_query") | |
| if selected_option == "Query all data": | |
| # If "Query all data" is selected, return the first class as a default for now | |
| # You can modify this behavior as needed | |
| return existing_classes[0] | |
| else: | |
| return selected_option | |
| def handle_new_class_selection(client, selected_class): | |
| if selected_class == "New Class": | |
| class_name = st.text_input("Enter the new class name:") | |
| class_description = st.text_input("Enter a description for the class:") | |
| if class_name and class_description: | |
| if st.button("Create Vector DB Class"): | |
| # Call function to create new class schema in Weaviate | |
| weaviate_utils.create_new_class_schema(client, class_name, class_description) | |
| def csv_upload_and_ingestion(client, selected_class): | |
| csv_file = st.file_uploader("Upload a CSV file", type=["csv"], key="csv_uploader") | |
| if csv_file: | |
| if st.button("Confirm CSV upload"): | |
| # Call function to ingest CSV data into Weaviate | |
| dataframe = weaviate_utils.ingest_data_to_weaviate(client, csv_file, selected_class) # Updated this line | |
| if dataframe is not None: # Check if ingestion was successful | |
| # Display a preview of the ingested data | |
| st.write(f"Your CSV was successfully integrated into the vector database under the class '{selected_class}'") | |
| st.write(dataframe.head()) # Display the first few rows of the dataframe as a preview | |
| st.session_state.csv_uploaded = True # Set session state variable | |
| def display_query_input(client, selected_class, tokenizer, model): # Added parameters | |
| question = st.text_input("Enter your question:", key="query_input") | |
| # Display the "Submit Query" button if CSV has been uploaded | |
| if st.session_state.get("csv_uploaded", False): | |
| if st.button("Submit Query"): | |
| if question: # Check if the question input is not empty | |
| # Call function to query TAPAS with selected data and entered question | |
| query_tapas_with_weaviate_data(client, selected_class, question, tokenizer, model) | |
| else: | |
| st.warning("Please provide a text query in the 'Enter your question:' input box to proceed.") | |
| st.write(f"Selected class type: {type(selected_class)}, Value: {selected_class}") | |
| def query_tapas_with_weaviate_data(client, selected_class, question, tokenizer, model): | |
| # 1. Perform hybrid search | |
| data = weaviate_utils.hybrid_search_weaviate(client, selected_class, question) | |
| st.write(f"Data from Weaviate: {data}") # Logging data from Weaviate | |
| # 2. Convert the data to TAPAS format | |
| table = weaviate_utils.convert_to_tapas_format(data) | |
| st.write(f"Data for TAPAS: {table}") # Logging data from TAPAS Table | |
| # 3. Call TAPAS with the table and the question | |
| answers = tapas_utils.ask_llm_chunk(tokenizer, model, table, [question]) | |
| st.write(f"TAPAS answers: {answers}") # Logging data from TAPAS Answers | |
| # Display the answers | |
| for answer in answers: | |
| st.write(f"Answer: {answer}") | |