the stack v2 is coming

Author: ChenghaoMou

poetry.lock

@@ -37,6 +37,7 @@ seaborn = "^0.12.2"
 ruff = "^0.0.265"
 insegel = "^1.3.1"
 sphinx-autobuild = "^2021.3.14"
+google-cloud-storage = "^2.10.0"
 
 [build-system]
 requires = ["poetry-core"]

pyproject.toml

@@ -0,0 +1,17 @@
+import sys
+
+from google.cloud import storage
+
+bucket = sys.argv[1]
+prefix = sys.argv[2]
+client = storage.Client()
+b = client.get_bucket(bucket)
+
+blobs = b.list_blobs(prefix=prefix)
+seen = set()
+for blob in blobs:
+    prefix = blob.name.rsplit("/", 1)[0]
+    dir = f"gs://{b.name}/{prefix}"
+    if dir not in seen:
+        print(dir)
+        seen.add(dir)

text_dedup/bigcode/ablation_visualize.ipynb

@@ -343,14 +343,19 @@ def area(s):
 # region: Quality Control
 def process_cluster(cluster: List[Any], enabled: bool = False) -> List[Any]:
     if not enabled:
-        np.random.shuffle(cluster)
+        RNG.shuffle(cluster)
         return cluster[:1]
 
     cluster.sort(
         key=lambda x: (
-            -x[-1] if x[-1] is not None else 0.0,  # star_events_count
-            -x[-2] if x[-2] is not None else 0.0,  # fork_events_count
-            -np.datetime64(x[-3]).astype(np.uint64) if x[-3] is not None else 0.0,  # visit_date
+            # license_type, the more permissive the better
+            ["permissive", "no_license", "non_permissive"].index(x[-1]) if x[-1] is not None else float("inf"),
+            # star_events_count, the more the better
+            -x[-2] if x[-2] is not None else 0.0,
+            # fork_events_count, the more the better
+            -x[-3] if x[-3] is not None else 0.0,
+            # visit_date, the earliest the better, tie breaker
+            np.datetime64(x[-4]).astype(np.uint64) if x[-4] is not None else float("inf"),
         )
     )
     return cluster[:1]
@@ -382,7 +387,7 @@ def partitioned_save(df: DataFrame, chunk_size: int, max_partitions: int, output
     """
 
     total_rows = df.count()
-    partitions = max(1, min(math.ceil(total_rows / chunk_size), max_partitions))
+    partitions = max(256, min(math.ceil(total_rows / chunk_size), max_partitions))
 
     def save_partition(df: pd.DataFrame) -> pd.DataFrame:  # type: ignore
         pid = df["__pid__"].iloc[0]
@@ -392,15 +397,17 @@ def save_partition(df: pd.DataFrame) -> pd.DataFrame:  # type: ignore
         )
         return pd.DataFrame([{"__status__": True, "__pid__": pid}])
 
+    log.debug(f"Saving {total_rows} rows to {partitions} partitions.")
+
     results = (
         df.repartition(partitions)  # random and uniform hash partitioning
         .withColumn("__pid__", F.spark_partition_id())
         .groupBy("__pid__")
         .applyInPandas(save_partition, schema="__status__ boolean, __pid__ int")
-        .toPandas()
+        .cache()
     )
 
-    if results["__status__"].all():
+    if results.filter(~F.col("__status__")).count() == 0:
         pd.DataFrame([]).to_csv(os.path.join(output, "_SUCCESS"), index=False, header=False)
         return
 
@@ -433,6 +440,8 @@ def save_partition(df: pd.DataFrame) -> pd.DataFrame:  # type: ignore
     conf = SparkConf()
     conf.set("spark.app.name", "MinHashLSH")
     conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
+    conf.set("spark.storage.memoryFraction", "1")
+    conf.set("spark.default.parallelism", "100")
     spark = SparkSession.builder.config(conf=conf).getOrCreate()  # type: ignore
     log: Logger = spark.sparkContext._jvm.org.apache.log4j.LogManager.getLogger(__name__)  # type: ignore
 
@@ -445,8 +454,8 @@ def save_partition(df: pd.DataFrame) -> pd.DataFrame:  # type: ignore
     if B is None or R is None:
         B, R = optimal_param(args.threshold, args.num_perm)
 
-    MAX_WRITE_CHUNK_SIZE: int = 1_000_000
-    MAX_WRITE_PARTITIONS: int = 256
+    MAX_WRITE_CHUNK_SIZE: int = 80_000
+    MAX_WRITE_PARTITIONS: int = 2048
     HASH_RANGES: List[Tuple[int, int]] = [(i * R, (i + 1) * R) for i in range(B)]
     PERMUTATIONS: Tuple[np.ndarray, np.ndarray] = (
         RNG.randint(1, MOD_PRIME, size=(args.num_perm,), dtype=DTYPE),
@@ -455,6 +464,7 @@ def save_partition(df: pd.DataFrame) -> pd.DataFrame:  # type: ignore
 
     # region: Data Loading
     df: DataFrame = spark.read.option("mergeSchema", "true").parquet(args.input)
+    # df = df.filter(F.col("license_type") == "permissive").cache()
     if args.index_column is None:
         df = df.withColumn("__id__", F.monotonically_increasing_id()).cache()
     else:
@@ -600,11 +610,9 @@ def save_partition(df: pd.DataFrame) -> pd.DataFrame:  # type: ignore
             "__component__",
             args.repo_column,
             "visit_date",
+            "fork_events_count",
             "star_events_count",
-            "fork_events_count"
-            # "max_stars_repo_stars_event_min_datetime",
-            # "max_stars_count",
-            # "max_forks_count",
+            "license_type",
         ]
         if args.rank
         else [

text_dedup/bigcode/dirs.list

@@ -6,9 +6,9 @@
 CLUSTER_NAME="chenghao-temp"
 PROJECT_ID="huggingface-science-codeparrot"
 REGION="us-central1"
-CONTAINER=""
-DIRECTORY=""
-NUM_WORKERS=18
+CONTAINER="gs://the_stack_v2"
+DIRECTORY="licensed_files"
+NUM_WORKERS=25
 MASTER_MACHINE_TYPE="c2d-standard-16"
 MASTER_BOOT_DISK_SIZE=1024
 WORKER_MACHINE_TYPE="c2-standard-16"
@@ -19,8 +19,9 @@ THRESHOLD=0.7
 REPO_COLUMN="repo_url"
 
 DEDUPED_DIRECTORY="${DIRECTORY}_deduped"
-DEDUPED_INDEX_DIRECTORY="${DEDUPED_DIRECTORY}_index"
-DIRS=$(gsutil ls "${CONTAINER}/${DIRECTORY}")
+# DEDUPED_INDEX_DIRECTORY="${DEDUPED_DIRECTORY}_index"
+# DIRS=("gs://the_stack_v2/licensed_files/language_id=Python")
+DIRS=$(cat dirs.list)
 
 # Create cluster if it doesn't exist
 if ! gcloud dataproc clusters list --region $REGION | grep -q $CLUSTER_NAME; then
@@ -49,12 +50,11 @@ i=0
 
 echo "Total number of directories: $TOTAL"
 for DIR in $DIRS; do
-
     # Progress bar
     echo -n "[ "
     curr_pos=$((i * LENGTH / TOTAL))
-    for ((k = 0 ; k <= curr_pos; k++)); do echo -n "==="; done
-    for ((j = k + 1; j <= LENGTH ; j++)); do echo -n "   "; done
+    for ((k = 0; k <= curr_pos; k++)); do echo -n "==="; done
+    for ((j = k + 1; j <= LENGTH; j++)); do echo -n "   "; done
     v=$(((i + 1) * 100 / TOTAL))
     echo -n " ] "
     echo "$v %" $'\r'
@@ -64,8 +64,8 @@ for DIR in $DIRS; do
     INPUT_GCS_PATH="${DIR}"
     LAN=$(echo "$DIR" | rev | cut -d'/' -f1 | rev)
     OUTPUT_GCS_PATH="${CONTAINER}/${DEDUPED_DIRECTORY}/${LAN}"
-    OUTPUT_INDEX_GCS_PATH="${CONTAINER}/${DEDUPED_INDEX_DIRECTORY}/${LAN}"
-    OUTPUT_STATUS_GCS_PATH="${CONTAINER}/${DEDUPED_INDEX_DIRECTORY}/${LAN}/_SUCCESS"
+    # OUTPUT_INDEX_GCS_PATH="${CONTAINER}/${DEDUPED_INDEX_DIRECTORY}/${LAN}"
+    OUTPUT_STATUS_GCS_PATH="${OUTPUT_GCS_PATH}/_SUCCESS"
     result=$(gsutil stat "${OUTPUT_STATUS_GCS_PATH}" 2>&1 | grep -c "No URLs matched")
     if [[ $result != 1 ]]; then
         echo "Skipping ${LAN}"
@@ -82,10 +82,13 @@ for DIR in $DIRS; do
         --input "$INPUT_GCS_PATH" \
         --output "$OUTPUT_GCS_PATH" \
         --threshold $THRESHOLD \
-        --output_index "$OUTPUT_INDEX_GCS_PATH" \
         --repo_column $REPO_COLUMN \
-        --rank \
-        --debug
+        --rank
+
+    # --rank
+    # --debug
+    # --output_index "$OUTPUT_INDEX_GCS_PATH"
+
 done
 
 gcloud dataproc clusters stop $CLUSTER_NAME --region $REGION