s2 dump scripts...

1 files changed, 30 insertions, 0 deletions

diff --git a/s2-dump-ids.py b/s2-dump-ids.py
new file mode 100644
index 00000000..66ff6d77
--- /dev/null
+++ b/s2-dump-ids.py
@@ -0,0 +1,30 @@
+import os
+import gzip
+import glob
+import json
+import click
+from util import *
+
+PAPER_JSON_DIR = 'datasets/s2/papers'
+
+@click.command()
+def s2_dump_ids():
+  ids = {}
+  for fn in glob.iglob('{}/**/*.json'.format(PAPER_JSON_DIR), recursive=True):
+    process_paper(fn, ids)
+  id_list = list(ids.keys())
+  print("Wrote {} ids".format(len(id_list)))
+  write_json('ids.json', id_list)
+
+def process_paper(fn, ids):
+  with open(fn, 'r') as f:
+    data = json.load(f)
+    ids[data['paperId']] = True
+    for cite in data['citations']:
+      ids[cite['paperId']] = True
+
+def paper_path(paper_id):
+  return '{}/{}/{}'.format(DATA_DIR, paper_id[0:3], paper_id)
+  
+if __name__ == '__main__':
+  s2_dump_ids()