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 10ed0bc0..2682e576 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 = ''