Upload 4 files

evaluation/README.md

@@ -0,0 +1,8 @@
+
+### Commands for running evaluation
+
+```console
+python evaluate.py --eval-dataset doc2dial
+python evaluate.py --eval-dataset quac
+python evaluate.py --eval-dataset qrecc
+```

evaluation/arguments.py

@@ -0,0 +1,31 @@
+
+import argparse
+import os
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Dragon-multiturn")
+
+    parser.add_argument('--query-encoder-path', type=str, default='nvidia/dragon-multiturn-query-encoder')
+    parser.add_argument('--context-encoder-path', type=str, default='nvidia/dragon-multiturn-context-encoder')
+
+    parser.add_argument('--data-folder', type=str, default='', help='path to the datafolder of ChatRAG Bench')
+    parser.add_argument('--eval-dataset', type=str, default='', help='evaluation dataset (e.g., doc2dial)')
+
+    parser.add_argument('--doc2dial-datapath', type=str, default='doc2dial/test.json')
+    parser.add_argument('--doc2dial-docpath', type=str, default='doc2dial/documents.json')
+
+    parser.add_argument('--quac-datapath', type=str, default='quac/test.json')
+    parser.add_argument('--quac-docpath', type=str, default='quac/documents.json')
+    
+    parser.add_argument('--qrecc-datapath', type=str, default='qrecc/test.json')
+    parser.add_argument('--qrecc-docpath', type=str, default='qrecc/documents.json')
+    
+    parser.add_argument('--topiocqa-datapath', type=str, default='topiocqa/dev.json')
+    parser.add_argument('--topiocqa-docpath', type=str, default='')
+    
+    parser.add_argument('--inscit-datapath', type=str, default='inscit/dev.json')
+    parser.add_argument('--inscit-docpath', type=str, default='')
+
+    args = parser.parse_args()
+
+    return args

evaluation/dataset.py

@@ -0,0 +1,74 @@
+
+import json
+
+
+def get_query(messages, num_turns=5):
+    ## convert query into a format as follows:
+    ## user: {user}\nagent: {agent}\nuser: {user}
+    query = ""
+    for item in messages[-num_turns:]:
+        item['role'] = item['role'].replace("assistant", "agent")
+        query += "{}: {}\n".format(item['role'], item['content'])
+    query = query.strip()
+    
+    return query
+
+
+def get_query_with_topic(messages, topic, num_turns=3):
+    ## convert query into a format as follows:
+    ## user: this is a question about {topic}. {user}\nagent: {agent}\nuser: this is a question about {topic}. {user}
+    query = ""
+    for item in messages[-num_turns:]:
+        item['role'] = item['role'].replace("assistant", "agent")
+        if item['role'] == 'user':
+            query += "{}: this is a question about {}. {}\n".format(item['role'], topic, item['content'])
+        else:
+            query += "{}: {}\n".format(item['role'], item['content'])
+    query = query.strip()
+
+    return query
+
+
+def get_data_for_evaluation(input_datapath, document_datapath, dataset_name):
+
+    print('reading evaluation data from %s' % input_datapath)
+    with open(input_datapath, "r") as f:
+        input_list = json.load(f)
+    
+    print('reading documents from %s' % document_datapath)
+    with open(document_datapath, "r") as f:
+        documents = json.load(f)
+
+    eval_data = {}
+    for item in input_list:
+        """
+        We incorporate topic information for topiocqa and inscit datasets:
+        query = get_query_with_topic(item['messages'], item['topic'])
+        """
+        query = get_query(item['messages'])
+
+        doc_id = item['document']
+        gold_idx = item['ground_truth_ctx']['index']
+
+        if dataset_name == 'qrecc':
+            """
+            The 'gold context' for the qrecc dataset is obtained based on the word
+            overlaps between gold answer and each context in the document, which might
+            not be the real gold context.
+
+            To improve the evaluation quality of this dataset,
+            we further add the answer of the query into the 'gold context' 
+            to ensure the 'gold context' is the most relevant chunk to the query.
+            
+            Note that this is just for the retrieval evaluation purpose, we do not
+            add answer to the context for the ChatRAG evaluation.
+            """
+            answer = item['answers'][0]
+            documents[doc_id][gold_idx] += " || "  + answer
+        
+        if doc_id not in eval_data:
+            eval_data[doc_id] = [{"query": query, "gold_idx": gold_idx}]
+        else:
+            eval_data[doc_id].append({"query": query, "gold_idx": gold_idx})
+
+    return eval_data, documents

evaluation/evaluate.py

@@ -0,0 +1,102 @@
+
+from transformers import AutoModel, AutoTokenizer
+from dataset import get_data_for_evaluation
+from arguments import get_args
+from tqdm import tqdm
+import torch
+import os
+
+
+def run_retrieval(eval_data, documents, query_encoder, context_encoder, tokenizer, max_seq_len=512):
+    
+    ranked_indices_list = []
+    gold_index_list = []
+    for doc_id in tqdm(eval_data):
+        context_list = documents[doc_id]
+        
+        with torch.no_grad():
+            # get chunk embeddings
+            context_embs = []
+            for chunk in context_list:
+                chunk_ids = tokenizer(chunk, max_length=max_seq_len, truncation=True, return_tensors="pt").to("cuda")
+
+                c_emb = context_encoder(input_ids=chunk_ids.input_ids, attention_mask=chunk_ids.attention_mask)
+                c_emb = c_emb.last_hidden_state[:, 0, :]
+                context_embs.append(c_emb)
+            context_embs = torch.cat(context_embs, dim=0)   # (num_chunk, hidden_dim)
+    
+            sample_list = eval_data[doc_id]
+            query_embs = []
+            for item in sample_list:
+                gold_idx = item['gold_idx']
+                gold_index_list.append(gold_idx)
+
+                query = item['query']
+                query_ids = tokenizer(query, max_length=max_seq_len, truncation=True, return_tensors="pt").to("cuda")
+                q_emb = query_encoder(input_ids=query_ids.input_ids, attention_mask=query_ids.attention_mask)
+                q_emb = q_emb.last_hidden_state[:, 0, :]
+                query_embs.append(q_emb)
+            query_embs = torch.cat(query_embs, dim=0)   # (num_query, hidden_dim)
+
+            similarities = query_embs.matmul(context_embs.transpose(0,1))     # (num_query, num_chunk)
+            ranked_results = torch.argsort(similarities, dim=-1, descending=True)   # (num_query, num_chunk)
+            ranked_indices_list.extend(ranked_results.tolist())
+
+    return ranked_indices_list, gold_index_list
+
+
+def calculate_recall(ranked_indices_list, gold_index_list, topk):
+    hit = 0
+    for ranked_indices, gold_index in zip(ranked_indices_list, gold_index_list):
+        for idx in ranked_indices[:topk]:
+            if idx == gold_index:
+                hit += 1
+                break
+    recall = hit / len(ranked_indices_list)
+
+    print("top-%d recall score: %.4f" % (topk, recall))
+
+
+def main():
+    args = get_args()
+
+    ## get tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(args.query_encoder_path)
+
+    ## get retriever model
+    query_encoder = AutoModel.from_pretrained(args.query_encoder_path)
+    context_encoder = AutoModel.from_pretrained(args.context_encoder_path)
+    query_encoder.to("cuda"), query_encoder.eval()
+    context_encoder.to("cuda"), context_encoder.eval()
+
+    ## get evaluation data
+    if args.eval_dataset == "doc2dial":
+        input_datapath = os.path.join(args.data_folder, args.doc2dial_datapath)
+        input_docpath = os.path.join(args.data_folder, args.doc2dial_docpath)
+    elif args.eval_dataset == "quac":
+        input_datapath = os.path.join(args.data_folder, args.quac_datapath)
+        input_docpath = os.path.join(args.data_folder, args.quac_docpath)
+    elif args.eval_dataset == "qrecc":
+        input_datapath = os.path.join(args.data_folder, args.qrecc_datapath)
+        input_docpath = os.path.join(args.data_folder, args.qrecc_docpath)
+    elif args.eval_dataset == "topiocqa" or args.eval_dataset == "inscit":
+        raise Exception("We have prepare the function to get queries, but a wikipedia corpus needs to be downloaded")
+    else:
+        raise Exception("Please input a correct eval_dataset name!")
+
+    eval_data, documents = get_data_for_evaluation(input_datapath, input_docpath, args.eval_dataset)
+    
+    ## run retrieval
+    ranked_indices_list, gold_index_list = run_retrieval(eval_data, documents, query_encoder, context_encoder, tokenizer)
+    print("number of the total test samples: %d" % len(ranked_indices_list))
+
+    ## calculate recall scores
+    print("evaluating on %s" % args.eval_dataset)
+    topk_list = [1, 5, 20]
+    for topk in topk_list:
+        calculate_recall(ranked_indices_list, gold_index_list, topk=topk)
+
+
+if __name__ == "__main__":
+    main()
+