File size: 3,708 Bytes
d934c26
f419ec9
168f49e
d934c26
 
7ba8571
 
 
 
 
d934c26
 
 
 
 
 
 
 
 
 
 
 
f419ec9
7ba8571
d934c26
 
f419ec9
d934c26
b9dd2dd
d934c26
 
 
 
 
 
f419ec9
d934c26
b9dd2dd
f419ec9
dffd9dd
 
 
 
 
 
 
 
 
 
e26e61e
 
d867f48
 
382ab41
 
d934c26
382ab41
 
 
 
 
 
d867f48
168f49e
 
 
 
 
 
 
 
 
 
 
 
 
 
b10208c
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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"]
        return st.selectbox("Select a class or query all data:", class_options, key="class_selector_query")

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.")



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)

    # 2. Convert the data to TAPAS format
    table = weaviate_utils.convert_to_tapas_format(data)

    # 3. Call TAPAS with the table and the question
    answers = tapas_utils.ask_llm_chunk(tokenizer, model, table, [question])

    # Display the answers
    for answer in answers:
        st.write(f"Answer: {answer}")