diff options
Diffstat (limited to 'gnqa/paper2_eval/src/document_operations.py')
| -rw-r--r-- | gnqa/paper2_eval/src/document_operations.py | 36 |
1 files changed, 36 insertions, 0 deletions
diff --git a/gnqa/paper2_eval/src/document_operations.py b/gnqa/paper2_eval/src/document_operations.py index 10ed0bc..2682e57 100644 --- a/gnqa/paper2_eval/src/document_operations.py +++ b/gnqa/paper2_eval/src/document_operations.py @@ -8,6 +8,15 @@ from r2r import R2R, Document, GenerationConfig, R2RClient class DocOps: _type = '' + values_key = { + "text" : {"name": "contexts", "append": 1}, + "associatedQuery": {"name": "question", "append": 0}, + "id": {"name": "id", "append": 1}, + "title": {"name": "titles", "append": 1}, + "document_id": {"name": "document_id", "append": 1}, + "extraction_id": {"name": "extraction_id", "append": 1}, + "content": {"name": "answer", "append": 0} + } def __init__(self): self._type = 'QuestionList' @@ -22,6 +31,33 @@ class DocOps: with open(outp_file, "a") as the_data: the_data.write(output) + def get_ragas_out_dict(): + return { "titles": [], + "extraction_id": [], + "document_id": [], + "id": [], + "contexts": [], + "answer": "", + "question": ""} + + + def extract_response(obj, values_key, thedict): + if isinstance(obj, dict): + for key, val in obj.items(): + if (key in values_key.keys()): + if (values_key[key]["append"]): + thedict[values_key[key]["name"]].append(val.replace("\n", " ").strip()) + else: + thedict[values_key[key]["name"]] = val.replace("\n", " ").strip() + print(("", "Key -> {0}\tValue -> {1}".format(key,val)) [verbose]) + else: + if (len(obj.items()) == 1 ): + print(key, " --> ", val) + extract_response(val, values_key, thedict) + elif isinstance(obj, list): + for item in obj: + extract_response(item, values_key, thedict) + class QuestionList: _verbose = 0 _doc = '' |