Fixes for correct archiving with low RAM

Author: NyanNyanovich

archive.sh

@@ -1,17 +1,28 @@
+#!/bin/bash
+set -euo pipefail
+
+echo "Exporting documents"
 python3 -m scripts.mongo_to_jsonl --output-path data/raw_new_docs.jsonl
+
+echo "Cleaning documents"
 python3 -m scripts.clean_docs --input-path data/raw_new_docs.jsonl --output-path data/new_docs.jsonl
 python3 -m scripts.clean_docs --input-path data/raw_old_docs.jsonl --output-path data/old_docs.jsonl
 
 cp data/old_docs.jsonl data/all_docs.jsonl
 cat data/new_docs.jsonl >> data/all_docs.jsonl
 
+echo "Filtering documents"
 python3 -m scripts.filter_documents data/all_docs.jsonl data/documents.jsonl
 
+echo "Exporting clusters"
 python3 -m scripts.clusters_to_jsonl --output-path data/new_clusters.jsonl
 cp data/old_clusters.jsonl data/all_clusters.jsonl
 cat data/new_clusters.jsonl >> data/all_clusters.jsonl
+
+echo "Filtering clusters"
 python3 -m scripts.filter_posted_clusers data/all_clusters.jsonl data/clusters.jsonl data/documents.jsonl
 
+echo "Packing"
 cp channels.json data/channels.json
 rm data/nyan_archive.tar.gz
 cd data && tar -czvf nyan_archive.tar.gz clusters.jsonl documents.jsonl channels.json

scripts/clusters_to_jsonl.py

@@ -8,7 +8,10 @@
 
 def main(
     output_path,
-    mongo_config
+    mongo_config,
+    clid_start,
+    clid_end,
+    batch_size
 ):
     with open(mongo_config) as r:
         config = json.load(r)
@@ -18,19 +21,36 @@ def main(
     clusters_collection_name = config["clusters_collection_name"]
     collection = client[database_name][clusters_collection_name]
 
-    clusters = list(collection.find({}))
-    clusters.sort(key=lambda x: x["annotation_doc"]["pub_time"])
+    if not clid_start:
+        first_cluster = collection.find_one(sort=[("clid", 1)])
+        clid_start = first_cluster["clid"]
+    if not clid_end:
+        last_cluster = collection.find_one(sort=[("clid", -1)])
+        clid_end = last_cluster["clid"] + 1
+    print(f"Start clid: {clid_start}")
+    print(f"End clid: {clid_end}")
+
+    current_clid_start = clid_start
     with open(output_path, "w") as w:
-        for cluster in clusters:
-            cluster.pop("_id")
-            cluster["annotation_doc"].pop("embedding", None)
-            cluster["annotation_doc"].pop("embedded_images", None)
-            w.write(json.dumps(cluster, ensure_ascii=False) + "\n")
+        while current_clid_start < clid_end:
+            print(clid_end - current_clid_start)
+            current_clid_end = current_clid_start + batch_size
+            clusters = list(collection.find({"clid": {"$gte": current_clid_start, "$lt": current_clid_end}}))
+            clusters.sort(key=lambda x: x["annotation_doc"]["pub_time"])
+            for cluster in clusters:
+                cluster.pop("_id")
+                cluster["annotation_doc"].pop("embedding", None)
+                cluster["annotation_doc"].pop("embedded_images", None)
+                w.write(json.dumps(cluster, ensure_ascii=False) + "\n")
+            current_clid_start = current_clid_end
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--output-path", type=str, default="data/clusters.jsonl")
     parser.add_argument("--mongo-config", type=str, default="configs/mongo_config.json")
+    parser.add_argument("--clid-start", type=int, default=None)
+    parser.add_argument("--clid-end", type=int, default=None)
+    parser.add_argument("--batch-size", type=int, default=1000)
     args = parser.parse_args()
     main(**vars(args))

scripts/filter_documents.py

@@ -1,16 +1,32 @@
 import sys
 import json
 from collections import Counter
+from tqdm import tqdm
 
 input_path = sys.argv[1]
 output_path = sys.argv[2]
 
-with open(input_path) as r, open(output_path, "w") as w:
-    documents = [json.loads(line) for line in r]
-    documents.sort(key=lambda x: x["pub_time"])
+skip_idx = set()
+with open(input_path) as r:
+    records = []
+    for idx, line in enumerate(tqdm(r)):
+        record = json.loads(line)
+        record = {k: v for k, v in record.items() if k in ("pub_time", "url")}
+        record["idx"] = idx
+        records.append(record)
+    records.sort(key=lambda x: x["pub_time"])
     used_urls = set()
-    for doc in documents:
+    for doc in tqdm(records):
         if doc["url"] in used_urls:
+            skip_idx.add(doc["idx"])
             continue
         used_urls.add(doc["url"])
+print("Found {} duplicates".format(len(skip_idx)))
+
+
+with open(input_path) as r, open(output_path, "w") as w:
+    for idx, line in enumerate(tqdm(r)):
+        if idx in skip_idx:
+            continue
+        doc = json.loads(line)
         w.write(json.dumps(doc, ensure_ascii=False).strip() + "\n")

scripts/filter_posted_clusers.py

@@ -9,51 +9,34 @@
 docs_path = sys.argv[3]
 
 with open(input_path) as r, open(output_path, "w") as w, open(docs_path, "r") as df:
-    docs = [json.loads(line) for line in df]
-    url2doc = {doc["url"]: doc for doc in docs}
-    clusters = [json.loads(line) for line in r]
-
-    filtered_clusters = []
+    urls = {json.loads(line)["url"] for line in df}
     seen = set()
-    for cluster in clusters:
+    for line in r:
+        cluster = json.loads(line)
         if isinstance(cluster["annotation_doc"], str):
             url = cluster["annotation_doc"]
-            if url not in url2doc:
+            if url not in urls:
                 continue
-            cluster["annotation_doc"] = url2doc[url]
-
             url = cluster["first_doc"]
-            if url not in url2doc:
+            if url not in urls:
                 continue
-            cluster["first_doc"] = url2doc[url]
-
-            cluster["docs"] = [url2doc[url] for url in cluster["docs"] if url in url2doc]
-
-        url = cluster["annotation_doc"]["url"]
-        if url not in url2doc or url in seen:
-            continue
-        seen.add(url)
-        filtered_clusters.append(cluster)
-    clusters = filtered_clusters
+            cluster["docs"] = [url for url in cluster["docs"] if url in urls]
+        else:
+            annot_doc = cluster["annotation_doc"]
+            if annot_doc["url"] not in urls:
+                continue
+            cluster["annotation_doc"] = annot_doc["url"]
 
-    clusters.sort(key=lambda x: x["first_doc"]["pub_time"])
-    for cluster in tqdm(clusters):
-        fixed_docs = []
-        for doc in cluster["docs"]:
-            new_doc = url2doc.get(doc["url"])
-            if not new_doc:
+            first_doc = cluster["first_doc"]
+            if first_doc["url"] not in urls:
                 continue
-            fixed_docs.append(doc["url"])
-        cluster["docs"] = fixed_docs
+            cluster["first_doc"] = first_doc["url"]
 
-        annot_doc = cluster["annotation_doc"]
-        if annot_doc["url"] not in url2doc:
-            continue
-        cluster["annotation_doc"] = annot_doc["url"]
+            fixed_docs = [doc["url"] for doc in cluster["docs"] if doc["url"] in urls]
+            cluster["docs"] = fixed_docs
 
-        first_doc = cluster["first_doc"]
-        if first_doc["url"] not in url2doc:
+        url = cluster["annotation_doc"]
+        if url in seen:
             continue
-        cluster["first_doc"] = first_doc["url"]
-
+        seen.add(url)
         w.write(json.dumps(cluster, ensure_ascii=False).strip() + "\n")

scripts/mongo_to_jsonl.py

@@ -9,7 +9,8 @@
 def main(
     output_path,
     mongo_config,
-    annotated
+    annotated,
+    ts_start
 ):
     with open(mongo_config) as r:
         config = json.load(r)
@@ -22,8 +23,9 @@ def main(
         documents_collection_name = config["documents_collection_name"]
     collection = client[database_name][documents_collection_name]
 
-    first_doc = collection.find_one(sort=[("pub_time", 1)])
-    ts_start = first_doc["pub_time"]
+    if not ts_start:
+        first_doc = collection.find_one(sort=[("pub_time", 1)])
+        ts_start = first_doc["pub_time"]
     print(f"Start timestamp: {ts_start}")
     ts_end = get_current_ts()
     print(f"End timestamp: {ts_end}")
@@ -46,6 +48,7 @@ def main(
     parser = argparse.ArgumentParser()
     parser.add_argument("--output-path", type=str, default="data/docs.jsonl")
     parser.add_argument("--annotated", action="store_true")
+    parser.add_argument("--ts-start", type=int, default=None)
     parser.add_argument("--mongo-config", type=str, default="configs/mongo_config.json")
     args = parser.parse_args()
     main(**vars(args))