Update app.py

app.py

@@ -15,118 +15,139 @@ st.set_page_config(page_title="PSC Runtime",
                    page_icon='🌸', layout="centered")
 
 
-import torch
-fn = "dl19-gpt-3.5.pt"
-object = torch.load(fn)
 
-
-outputs = object[2]
-query2outputs = {}
-for output in outputs:
-    all_queries = {x['query'] for x in output}
-    assert len(all_queries) == 1
-    query = list(all_queries)[0]
-    query2outputs[query] = [x['hits'] for x in output]
-
-
-search_query = st.selectbox(
+name = st.selectbox(
     "",
-    sorted(query2outputs),
+    ["dl19", "dl20"]
     index=None,
-    placeholder="Choose a query from the list..."
+    placeholder="Choose a dataset..."
 )
 
-def preferences_from_hits(list_of_hits):
-    docid2id = {}
-    id2doc = {}
-    preferences = []
-
-    for result in list_of_hits:
-        for doc in result:
-            if doc["docid"] not in docid2id:
-                id = len(docid2id)
-                docid2id[doc["docid"]] = id
-                id2doc[id] = doc
-        print([doc["docid"] for doc in result])
-        print([docid2id[doc["docid"]] for doc in result])
-        preferences.append([docid2id[doc["docid"]] for doc in result])
-
-    #  = {v: k for k, v in docid2id.items()}
-    return np.array(preferences), id2doc
-
-
-def load_qrels(name):
-    import ir_datasets
-    if name == "dl19":
-        ds_name = "msmarco-passage/trec-dl-2019/judged"
-    elif name == "dl20":
-        ds_name = "msmarco-passage/trec-dl-2020/judged"
-    else:
-        raise ValueError(name)
-
-    dataset = ir_datasets.load(ds_name)
-    qrels = defaultdict(dict)
-    for qrel in dataset.qrels_iter():
-        qrels[qrel.query_id][qrel.doc_id] = qrel.relevance
-    return qrels
-
-
-def aggregate(list_of_hits):
-    import numpy as np
-    from permsc import KemenyOptimalAggregator, sum_kendall_tau, ranks_from_preferences
-    from permsc import BordaRankAggregator
-
-    preferences, id2doc = preferences_from_hits(list_of_hits)
-    y_optimal = KemenyOptimalAggregator().aggregate(preferences)
-    # y_optimal = BordaRankAggregator().aggregate(preferences)
-
-    return [id2doc[id] for id in y_optimal]
-
-
-def write_ranking(search_results):
-    # st.write(
-    #     f'<p align=\"right\" style=\"color:grey;\"> Before aggregation for query [{search_query}] ms</p>', unsafe_allow_html=True)
-
-    qid = {result["qid"] for result in search_results}
-    assert len(qid) == 1
-    qid = list(qid)[0]
+model_name = st.selectbox(
+    "",
+    ["gpt-3.5", "gpt-4"]
+    index=None,
+    placeholder="Choose a model..."
+)
 
-    for i, result in enumerate(search_results):
-        result_id = result["docid"]
-        contents = result["content"]
+# "dl19"
 
-        label = qrels[str(qid)].get(str(result_id), 0)
-        if label == 3:
-            style = "style=\"color:blue;\""
-        elif label == 2:
-            style = "style=\"color:green;\""
-        elif label == 1:
-            style = "style=\"color:red;\""
+if name and model_name:
+    
+    import torch
+    # fn = f"dl19-gpt-3.5.pt"
+    fn = f"{name}-{model_name}.pt"
+    object = torch.load(fn)
+    
+    
+    outputs = object[2]
+    query2outputs = {}
+    for output in outputs:
+        all_queries = {x['query'] for x in output}
+        assert len(all_queries) == 1
+        query = list(all_queries)[0]
+        query2outputs[query] = [x['hits'] for x in output]
+    
+    
+    search_query = st.selectbox(
+        "",
+        sorted(query2outputs),
+        index=None,
+        placeholder="Choose a query from the list..."
+    )
+    
+    def preferences_from_hits(list_of_hits):
+        docid2id = {}
+        id2doc = {}
+        preferences = []
+    
+        for result in list_of_hits:
+            for doc in result:
+                if doc["docid"] not in docid2id:
+                    id = len(docid2id)
+                    docid2id[doc["docid"]] = id
+                    id2doc[id] = doc
+            print([doc["docid"] for doc in result])
+            print([docid2id[doc["docid"]] for doc in result])
+            preferences.append([docid2id[doc["docid"]] for doc in result])
+    
+        #  = {v: k for k, v in docid2id.items()}
+        return np.array(preferences), id2doc
+    
+    
+    def load_qrels(name):
+        import ir_datasets
+        if name == "dl19":
+            ds_name = "msmarco-passage/trec-dl-2019/judged"
+        elif name == "dl20":
+            ds_name = "msmarco-passage/trec-dl-2020/judged"
         else:
-            style = "style=\"color:grey;\""
-
-        print(qid, result_id, label, style)
-        # output = f'<div class="row"> <b>Rank</b>: {i+1} | <b>Document ID</b>: {result_id} | <b>Score</b>:{result_score:.2f}</div>'
-        output = f'<div class="row" {style}> <b>Rank</b>: {i+1} | <b>Document ID</b>: {result_id}'
-
-        try:
-            st.write(output, unsafe_allow_html=True)
-            st.write(
-                f'<div class="row" {style}>{contents}</div>', unsafe_allow_html=True)
-
-        except:
-            pass
-        st.write('---')
-
-aggregated_ranking = aggregate(query2outputs[search_query])
-qrels = load_qrels("dl19")
-col1, col2 = st.columns([5, 5])
-
-if search_query:
-    with col1:
-        if search_query or button_clicked:
-            write_ranking(search_results=query2outputs[search_query][0])
-    
-    with col2:
-        if search_query or button_clicked:
-            write_ranking(search_results=aggregated_ranking)
+            raise ValueError(name)
+    
+        dataset = ir_datasets.load(ds_name)
+        qrels = defaultdict(dict)
+        for qrel in dataset.qrels_iter():
+            qrels[qrel.query_id][qrel.doc_id] = qrel.relevance
+        return qrels
+    
+    
+    def aggregate(list_of_hits):
+        import numpy as np
+        from permsc import KemenyOptimalAggregator, sum_kendall_tau, ranks_from_preferences
+        from permsc import BordaRankAggregator
+    
+        preferences, id2doc = preferences_from_hits(list_of_hits)
+        y_optimal = KemenyOptimalAggregator().aggregate(preferences)
+        # y_optimal = BordaRankAggregator().aggregate(preferences)
+    
+        return [id2doc[id] for id in y_optimal]
+    
+    
+    def write_ranking(search_results):
+        # st.write(
+        #     f'<p align=\"right\" style=\"color:grey;\"> Before aggregation for query [{search_query}] ms</p>', unsafe_allow_html=True)
+    
+        qid = {result["qid"] for result in search_results}
+        assert len(qid) == 1
+        qid = list(qid)[0]
+    
+        for i, result in enumerate(search_results):
+            result_id = result["docid"]
+            contents = result["content"]
+    
+            label = qrels[str(qid)].get(str(result_id), 0)
+            if label == 3:
+                style = "style=\"color:blue;\""
+            elif label == 2:
+                style = "style=\"color:green;\""
+            elif label == 1:
+                style = "style=\"color:red;\""
+            else:
+                style = "style=\"color:grey;\""
+    
+            print(qid, result_id, label, style)
+            # output = f'<div class="row"> <b>Rank</b>: {i+1} | <b>Document ID</b>: {result_id} | <b>Score</b>:{result_score:.2f}</div>'
+            output = f'<div class="row" {style}> <b>Rank</b>: {i+1} | <b>Document ID</b>: {result_id}'
+    
+            try:
+                st.write(output, unsafe_allow_html=True)
+                st.write(
+                    f'<div class="row" {style}>{contents}</div>', unsafe_allow_html=True)
+    
+            except:
+                pass
+            st.write('---')
+    
+    
+    aggregated_ranking = aggregate(query2outputs[search_query])
+    qrels = load_qrels(name)
+    col1, col2 = st.columns([5, 5])
+    
+    if search_query:
+        with col1:
+            if search_query or button_clicked:
+                write_ranking(search_results=query2outputs[search_query][0])
+        
+        with col2:
+            if search_query or button_clicked:
+                write_ranking(search_results=aggregated_ranking)