cc warc

levkk · levkk · commit aa88ecca16b3 · 2024-06-18T10:56:01.000-07:00
diff --git a/pgml-apps/cc-warc/.gitignore b/pgml-apps/cc-warc/.gitignore
@@ -0,0 +1,2 @@
+venv/
+__pycache__
diff --git a/pgml-apps/cc-warc/api.py b/pgml-apps/cc-warc/api.py
@@ -0,0 +1 @@
+import Flask
diff --git a/pgml-apps/cc-warc/cli.py b/pgml-apps/cc-warc/cli.py
@@ -0,0 +1,45 @@
+from warcio.archiveiterator import ArchiveIterator
+import click
+import asyncio
+import requests
+from pgml import Collection, Pipeline
+from config import DATABASE_URL
+
+collection = Collection("warc", DATABASE_URL)
+pipeline = Pipeline("warc_search", {
+	"body": {
+		"splitter": {"model": "recursive_character"},
+		"semantic_search": {
+			"model": "mixedbread-ai/mxbai-embed-large-v1",
+		}
+	}
+})
+
+async def ingest(paths, limit=500):
+	await collection.add_pipeline(pipeline)
+	with open(paths) as f:
+		for path in f:
+			req = requests.get("https://data.commoncrawl.org/%s" % path.strip(), stream=True)
+			batch = []
+			for record in ArchiveIterator(req.raw, arc2warc=True):
+				document = {
+					"id": record.rec_headers.get_header("WARC-Target-URI"),
+					"body": record.content_stream().read().decode("utf-8")
+				}
+				print(document)
+				batch.append(document)
+				if len(batch) == limit:
+					exit(1)
+					# await collection.upsert_documents(batch)
+					batch = []
+
+
+@click.command()
+@click.option("--path", help="Path to the WET paths file.", default="paths.txt")
+@click.option("--limit", default=5, help="How many files to download and ingest.")
+def cli(path, limit):
+	asyncio.run(ingest(path, limit))
+
+
+if __name__ == "__main__":
+	cli()
diff --git a/pgml-apps/cc-warc/config.py b/pgml-apps/cc-warc/config.py
@@ -0,0 +1,3 @@
+import os
+
+DATABASE_URL = os.environ.get("PGML_DATABASE_URL", "postgres:///pgml")
diff --git a/pgml-apps/cc-warc/paths.txt b/pgml-apps/cc-warc/paths.txt
diff --git a/pgml-apps/cc-warc/qa.txt b/pgml-apps/cc-warc/qa.txt
@@ -0,0 +1 @@
+Exception: `id` must be a key in document
diff --git a/pgml-apps/cc-warc/requirements.txt b/pgml-apps/cc-warc/requirements.txt
@@ -0,0 +1,8 @@
+blinker==1.8.2
+click==8.1.7
+Flask==3.0.3
+itsdangerous==2.2.0
+Jinja2==3.1.4
+MarkupSafe==2.1.5
+pgml==1.1.0
+Werkzeug==3.0.3

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+import os`
	`2`	`+`
	`3`	`+DATABASE_URL = os.environ.get("PGML_DATABASE_URL", "postgres:///pgml")`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Exception: `id` must be a key in document