working

.gitignore

@@ -2,4 +2,5 @@ datasets/
 __pycache__/
 env/
 .ipynb_checkpoints/
-*.ipynb
+*.ipynb
+*.pyc

app.py

@@ -8,8 +8,6 @@ import copy
 import re
 import tqdm
 import numpy as np
-import plotly.express as px
-from sentence_transformers import SentenceTransformer
 
 import pandas as pd
 from nltk.corpus import stopwords
@@ -28,7 +26,7 @@ nltk.download('punkt')
 nltk.download('stopwords')
 
 
-from dataset_loading import load_local_corpus, load_local_queries, load_local_triples
+from dataset_loading import load_local_corpus, load_local_queries, load_local_triples, load_json
 
 
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
@@ -38,69 +36,6 @@ current_checkboxes = []
 query_input = None
 
 
-
-@st.cache_data
-def create_sim_chart(triples_map, triples1_name, queries1, corpus, instances_to_use, triples2_map, triples2_pos_pids, triples2_name, queries2):
-    print(f"Calculating cosine similarity for {len(triples_map)} instances")
-    e5_model = SentenceTransformer('intfloat/e5-large-v2', trust_remote_code=True)
-    all_triples = [triples_map[str(inst_num)][0] for inst_num in instances_to_use]
-    # breakpoint()
-    corpus_texts = [corpus[str(pos_pid)] for _, pos_pid, _ in all_triples]
-    neg_corpus_texts = [corpus[str(neg_pid)] for _, _, neg_pid in all_triples]
-    queries_for_embed = [queries[str(qid)] for qid, _, _ in all_triples]
-    corpus_embeddings = e5_model.encode(corpus_texts)
-    neg_corpus_embeddings = e5_model.encode(neg_corpus_texts)
-    query_embeddings = e5_model.encode(queries1)
-    cos_sim_pos = []
-    cos_sim_neg = []
-    for query_emb, pos_emb, neg_emb in zip(query_embeddings, corpus_embeddings, neg_corpus_embeddings):
-        cos_sim_pos.append(np.dot(query_emb, pos_emb) / (np.linalg.norm(query_emb) * np.linalg.norm(pos_emb)))
-        cos_sim_neg.append(np.dot(query_emb, neg_emb) / (np.linalg.norm(query_emb) * np.linalg.norm(neg_emb)))
-    df_sim = pd.DataFrame({"triples": triples1_name, "cosine_similarity_pos": cos_sim_pos, "cosine_similarity_neg": cos_sim_neg})
-
-    if triples2 is not None:
-        print(f"Calculating cosine similarity for {len(triples2_pos_pids)} instances for 2")
-        all_triples2 = [triples2_map[str(inst_num)][0] for inst_num in instances_to_use]
-        corpus_texts2 = [corpus[str(pos_pid)] for _, pos_pid, _ in all_triples2]
-        neg_corpus_texts2 = [corpus[str(neg_pid)] for _, _, neg_pid in all_triples2]
-        queries_for_embed2 = [queries[str(qid)] for qid, _, _ in all_triples2]
-        corpus_embeddings2 = e5_model.encode(corpus_texts2)
-        neg_corpus_embeddings2 = e5_model.encode(neg_corpus_texts2)
-        query_embeddings = e5_model.encode(queries2)
-        cos_sim_pos = []
-        cos_sim_neg = []
-        for query_emb, pos_emb, neg_emb in zip(query_embeddings, corpus_embeddings, neg_corpus_embeddings):
-            cos_sim_pos.append(np.dot(query_emb, pos_emb) / (np.linalg.norm(query_emb) * np.linalg.norm(pos_emb)))
-            cos_sim_neg.append(np.dot(query_emb, neg_emb) / (np.linalg.norm(query_emb) * np.linalg.norm(neg_emb)))
-        df_sim2 = pd.DataFrame({"triples": triples2_name, "cosine_similarity_pos": cos_sim_pos, "cosine_similarity_neg": cos_sim_neg})
-
-        df_sim = pd.concat([df_sim, df_sim2])
-        queries_to_return = list(zip(queries1, queries2))
-    else:
-        queries_to_return = queries1
-    
-    return df_len, df_sim, queries_to_return
-
-
-
-
-@st.cache_data
-def convert_df(df):
-    # IMPORTANT: Cache the conversion to prevent computation on every rerun
-    return df.to_json(orient='records', lines=True, path_or_buf=None).encode('utf-8')
-
-
-def create_histogram_relevant_docs(relevant_df):
-    # turn results into a dataframe and then plot
-    fig = px.histogram(relevant_df, x="relevant_docs")
-    # make it fit in one column
-    fig.update_layout(
-        height=400,
-        width=250
-    )
-    return fig
-
-
 def get_current_data():
     cur_query_data = []
     cur_query = query_input.replace("\n", "\\n")
@@ -135,63 +70,57 @@ def validate(config_option, file_loaded):
 
 with st.sidebar:
     st.title("Options")
-    # add a text input for naming triples 1
-    st.header("Name of triples 1")
-    triples1_name = st.text_input("Name of triples 1", key="triples1_name")
-    # add a text input for naming triples 2
-    st.header("Name of triples 2")
-    triples2_name = st.text_input("Name of triples 2", key="triples2_name")
-    st.header("Upload corpus")
-    corpus_file = st.file_uploader("Choose a file", key="corpus")
-    corpus = load_local_corpus(corpus_file)
-    st.header("Upload queries")
-    queries_file = st.file_uploader("Choose a file", key="queries")
-    queries = load_local_queries(queries_file)
-    st.header("Upload triples")
-    triples_file = st.file_uploader("Choose a file", key="triples")
-    triples = load_local_triples(triples_file)
-    if triples:
-        triples_pos_pids = [item[1] for item in triples]
-        triples_map = {item[1]: [] for item in triples}
-        query1_ids = [item[0] for item in triples]
-        for item in triples:
-            triples_map[item[1]].append(item)
-
-    triples2_file = st.file_uploader("Choose a second triples file for comparison", key="triples2")
-    triples2 = load_local_triples(triples2_file)
-    if triples2:
-        triples2_pos_pids = [item[1] for item in triples2]
-        triples2_map = {item[1]: [] for item in triples2}
-        for item in triples2:
-            triples2_map[item[1]].append(item)
-
-        # filter them by their intersection
-        triples_pos_pids = list(set(triples_pos_pids).intersection(triples2_pos_pids))
-        triples2_pos_pids = list(set(triples2_pos_pids).intersection(triples_pos_pids))
-        triples_map = {item[1]: triples_map[item[1]] for item in triples if item[1] in triples_pos_pids}
-        triples2_map = {item[1]: triples2_map[item[1]] for item in triples2 if item[1] in triples2_pos_pids}
-        query2_ids = [item[0] for item in triples2]
-        print(f"There are {len(triples_pos_pids)} triples in common between the two files")
-
-    ## make sure all qids in triples are in queries and write out a warning if not
-    if queries is not None and triples is not None:
-        missing_qids = set(query1_ids) - set(queries.keys())
-        if len(missing_qids) > 0:
-            st.warning(f"The following qids in triples are not in queries: {missing_qids}. \nPlease fix")
-
-        if triples2 is not None:
-            missing_qids = set(query2_ids) - set(queries.keys())
-            if len(missing_qids) > 0:
-                st.warning(f"The following qids in triples2 are not in queries: {missing_qids}. \nPlease fix")
+    use_default = st.checkbox("Use default data", value=False)
+    if use_default:
+        st.write("Using default data")
+        with open("example_output.json", "r") as fin:
+            data = json.load(fin)
+            ids = data["ids"]
+            prompts = data["input"]
+            # they can be longer in partial cases
+            outputs = data["parsed"]
+            ids = data["ids"][:len(outputs)]
+            prompts = prompts[:len(outputs)]
+            mapping = {id: (prompt, output) for id, prompt, output in zip(ids, prompts, outputs)}
+            data2 = None
+    else:
+        st.header("Input File")
+        input_file = st.file_uploader("Choose a file", key="data")
+        data = load_json(input_file)
+        if data is not None:
+            ids = data["ids"]
+            prompts = data["input"]
+            # they can be longer in partial cases
+            outputs = data["parsed"]
+            ids = data["ids"][:len(outputs)]
+            prompts = prompts[:len(outputs)]
+            mapping = {id: (prompt, output) for id, prompt, output in zip(ids, prompts, outputs)}
+
+        input_file2 = st.file_uploader("Choose a second file", key="data2")
+        data2 = load_json(input_file2)
+        if data2 is not None:
+            ids2 = data2["ids"]
+            prompts2 = data2["input"]
+            # they can be longer in partial cases
+            outputs2 = data2["parsed"]
+            ids2 = data2["ids"][:len(outputs2)]
+            prompts2 = prompts2[:len(outputs2)]
+            mapping2 = {id: (prompt, output) for id, prompt, output in zip(ids2, prompts2, outputs2)}
 
+        
+
+   
 col1, col2 = st.columns([1, 3], gap="large")
 
-if corpus is not None and queries is not None and triples is not None:
+if data is not None:
+    joint_ids = ids if data2 is None else list(set(ids2).intersection(ids))
+    # print(f"Not using ids {set(ids) - set(joint_ids)} and {set(ids2) - set(joint_ids)}")
+    
     with st.sidebar:
         st.success("All files uploaded")
 
     with col1:
-        set_of_cols =  set(triples_pos_pids) if triples2 is None else set(triples2_pos_pids).intersection(triples_pos_pids)
+        set_of_cols = joint_ids
         container_for_nav = st.container()
         name_of_columns = sorted([item for item in set_of_cols])
         instances_to_use = name_of_columns
@@ -220,144 +149,43 @@ if corpus is not None and queries is not None and triples is not None:
         selectbox_instance = container_for_nav.selectbox("Select instance by ID", ["Overview"] + name_of_columns, on_change=sync_from_drop, key="selectbox_instance")
         st.divider()  
 
-        # get average words lengths for generated queries
-        queries1 = [queries[str(triples_map[str(inst_num)][0][0])] for inst_num in instances_to_use]
-        triples1_query_len = [len(query.split()) for query in queries1]
-        df_len = pd.DataFrame({"triples": triples1_name, "query_len": triples1_query_len})
-
-        if triples2 is not None:
-            queries2 = [queries[str(triples2_map[str(inst_num)][0][0])] for inst_num in instances_to_use]
-            triples2_query_len = [len(query.split()) for query in queries2]
-            df_len = pd.concat([df_len, pd.DataFrame({"triples": triples2_name, "query_len": triples2_query_len})])
-        else:
-            queries2 = None
-            triples2_pos_pids = None
-            triples2_map = None
-            triples2_pos_pids = None
-
-        # now load E5-small-v2 and calculate the cosine similarity of query and docs
-        df_len, df_sim_og, queries_to_return = create_sim_chart(triples_map, triples1_name, queries1, corpus, instances_to_use, triples2_map, triples2_pos_pids, triples2_name, queries2)
-        df_sim = df_sim_og.copy()
         
-        # make a plotly chart with the lengths (with both triples if available)
-        st.subheader("Query Lengths")
-        fig = px.histogram(df_len, x="query_len", color="triples", marginal="box", title="Query Lengths", barmode="overlay", color_discrete_map={triples1_name: "blue", triples2_name: "red"})
-        fig.update_layout(
-            height=400,
-            width=250
-        )
-        st.plotly_chart(fig)
-
-        # make a plotly plot of the cosine similarities like the above
-        st.subheader("Cosine Similarity")
-        fig = px.histogram(df_sim, x="cosine_similarity_pos", color="triples", marginal="box", title="Cosine Similarity", barmode="overlay", color_discrete_map={triples1_name: "blue", triples2_name: "red"})
-        fig.update_layout(
-            height=400,
-            width=250
-        )
-        st.plotly_chart(fig)
-
-        # make a plotly plot of the cosine similarities like the above
-        df_sim["Diff"] = df_sim["cosine_similarity_pos"] - df_sim["cosine_similarity_neg"]
-        st.subheader("Cosine Similarity")
-        fig = px.histogram(df_sim, x="Diff", color="triples", marginal="box", title="Cosine Similarity Diff (Pos - Neg)", barmode="overlay", color_discrete_map={triples1_name: "blue", triples2_name: "red"})
-        fig.update_layout(
-            height=400,
-            width=250
-        )
-        st.plotly_chart(fig)
-
-        # make a checkbox that if checked will download a CSV of the queries to return
-        if st.checkbox("Download data as JSONL"):
-            df_of_queries_to_return = pd.DataFrame({
-                triples1_name: [item[0] for item in queries_to_return],
-                triples2_name: [item[1] for item in queries_to_return]
-            })
-            st.download_button(
-                label="Download data as JSONL",
-                data=convert_df(df_of_queries_to_return),
-                file_name='queries_to_return.jsonl',
-                mime='text/json',
-            )
-
-
     with col2:
         # get instance number
         inst_index = number_of_col
 
         if inst_index >= 0:
-            inst_num = instances_to_use[inst_index]
-            
-            st.markdown("<h1 style='text-align: center; color: black;text-decoration: underline;'>Editor</h1>", unsafe_allow_html=True)
-
-
-            
-            if triples2 is not None:
-                triples_1_col, triples_2_col = st.columns([1, 1], gap="small")
-
-                with triples_1_col:
-                    container = st.container()
-                    container.divider()
-                    container.subheader(f"Triples 1: {triples1_name}")
-                    container.divider()
-                    container.subheader(f"Query")
-                    all_triples = triples_map[str(inst_num)][0]
-                    query_id = all_triples[0]
-                    query_text = queries[str(query_id)].strip()
-                    query_input = container.markdown(f"QID: {inst_num}\n\n{query_text}")
-                    container.divider()
-                    doc_texts = [(pos, neg, corpus[pos], corpus[neg]) for qid, pos, neg in [all_triples]]
-                    
-                    for i, (pos, neg, pos_text, neg_text) in enumerate(doc_texts):
-                        container.markdown(f"## Positive")
-                        container.markdown(f"\n{extract_doc_text(pos_text)}", True)
-                        container.markdown(f"## Negative")
-                        container.markdown(f"\n{extract_doc_text(neg_text)}", True)
-
-                with triples_2_col:
-                    container = st.container()
-                    container.divider()
-                    container.subheader(f"Triples 2: {triples2_name}")
-                    container.divider()
-                    container.subheader(f"Query")
-
-                    # all_triples = triples_map[str(inst_num)][0]
-                    # index_of_inst_num = triples_pos_pids.index(inst_num)
-                    all_triples2 = triples2_map[str(inst_num)][0]
-
-                    query_text = queries[str(all_triples2[0])].strip()
-                    query_input = container.markdown(f"QID: {all_triples2[0]}\n\n{query_text}")
-                    container.divider()
-                    doc_texts = [(pos, neg, corpus[pos], corpus[neg]) for qid, pos, neg in [all_triples2]]
-                    
-                    for i, (pos, neg, pos_text, neg_text) in enumerate(doc_texts):
-                        container.markdown(f"## Positive")
-                        container.markdown(f"\n{extract_doc_text(pos_text)}", True)
-                        container.markdown(f"## Negative")
-                        container.markdown(f"\n{extract_doc_text(neg_text)}", True)
+            prompt, output = mapping[joint_ids[inst_index]]
+            if data2 is not None:
+                prompt2, output2 = mapping2[joint_ids[inst_index]]
+                col1_out, col2_out = st.columns([1, 1], gap="small")
+
+                with col1_out:
+                    st.markdown("<h2 style='text-align: center; color: black;'>Prompt</h2>", unsafe_allow_html=True)
+                    st.write(prompt)
+
+                    st.markdown("<h2 style='text-align: center; color: black;'>Output</h2>", unsafe_allow_html=True)
+                    st.json(output)
+
+                with col2_out:
+                    st.markdown("<h2 style='text-align: center; color: black;'>Prompt</h2>", unsafe_allow_html=True)
+                    st.write(prompt2)
+
+                    st.markdown("<h2 style='text-align: center; color: black;'>Output</h2>", unsafe_allow_html=True)
+                    st.json(output2)
+
 
 
             else:
-                container = st.container()
-                container.divider()
-                st.subheader(f"Triples: {triples1_name}")
+                st.markdown("<h1 style='text-align: center; color: black;text-decoration: underline;'>Editor</h1>", unsafe_allow_html=True)
+
+                st.markdown("<h2 style='text-align: center; color: black;'>Prompt</h2>", unsafe_allow_html=True)
+                st.write(prompt)
 
-                container.subheader(f"Query")
-                print(inst_num)
-                all_triples = triples_map[str(inst_num)][0]
-                print(all_triples)
-                query_text = queries[str(all_triples[0])].strip()
-                query_input = container.markdown(f"QID: {all_triples[0]}\n\n{query_text}")
-                container.divider()
-                doc_texts = [(pos, neg, corpus[pos], corpus[neg]) for qid, pos, neg in [all_triples]]
-                    
-                for i, (pos, neg, pos_text, neg_text) in enumerate(doc_texts):
-                    container.markdown(f"## Positive")
-                    container.markdown(f"\n{extract_doc_text(pos_text)}", True)
-                    container.markdown(f"## Negative")
-                    container.markdown(f"\n{extract_doc_text(neg_text)}", True)
+                st.markdown("<h2 style='text-align: center; color: black;'>Output</h2>", unsafe_allow_html=True)
+                st.json(output)
 
-            
+                
 
 
 
@@ -377,4 +205,4 @@ if corpus is not None and queries is not None and triples is not None:
     
 
 else:
-    st.warning("Please choose a dataset and upload a run file. If you chose \"custom\" be sure that you uploaded all files (queries, corpus, triples)")
+    st.warning("Please choose an output file from prompting and upload it")

dataset_loading.py

@@ -5,7 +5,6 @@ import pandas as pd
 from collections import defaultdict
 import json
 import copy
-import plotly.express as px
 
 
 
@@ -130,6 +129,12 @@ def load_jsonl(f):
     return did2text, sub_did2text
 
 
+@st.cache_data
+def load_json(f):
+    if f is None:
+        return None
+    return json.load(f)
+
 
 @st.cache_data(persist="disk")
 def get_dataset(dataset_name: str, input_fields_doc, input_fields_query):