Update knowledge_level_and_agent_type.py

RTXteam · May 15, 2024 · 4419b78 · 4419b78
1 parent 58bd791
commit 4419b78
Showing 1 changed file with 13 additions and 2 deletions.
diff --git a/knowledge_level_and_agent_type.py b/knowledge_level_and_agent_type.py
@@ -19,6 +19,14 @@
 import kg2_util
 
 
+def make_arg_parser():
+    arg_parser = argparse.ArgumentParser(description='kg2_jsonl_to_kgx_jsonl.py: converts a JSONL file in the KG2 format to a JSONL file in the KGX format')
+    arg_parser.add_argument('input_edges_file', type=str)
+    arg_parser.add_argument('output_file_location', type=str)
+    arg_parser.add_argument('output_file', type=str)
+    return arg_parser
+
+
 # Check that all sources are represented in the agent_type dictionary
 def check_sources(agent_type_dict, sources):
     for source in sources:
@@ -30,11 +38,14 @@ def check_sources(agent_type_dict, sources):
 def process_edges(input_edges_file, output_file_location, output_file, agent_type_map):
 
     print(f"Adding knowledge_level and agent_type to edges from {input_edges_file} to {output_file_location}/{output_file}")
-    print(f"Type: {type(agent_type_map)} /nAgent type map: {agent_type_map}")
+  #  print(f"Type: {type(agent_type_map)} /nAgent type map: {agent_type_map}")
 
     edge_ctr = 0
     edges_read_jsonlines_data = kg2_util.start_read_jsonlines(input_edges_file)
     edges = edges_read_jsonlines_data[0]
+
+    edges_info = kg2_util.create_kg2_jsonlines()
+    edges_output = edges_info[0]
 
     for edge_dict in edges:
         edge_ctr += 1
@@ -44,7 +55,7 @@ def process_edges(input_edges_file, output_file_location, output_file, agent_typ
         # check_sources(agent_type_map, sources)
         edge_dict['agent_type'] = 'not_provided'
         edge_dict['knowledge_level'] = 'not_provided'
-        output_file.write(edge_dict)
+        edges_info.write(edge_dict)
 
 
 if __name__ == '__main__':