daft-examples/examples/commoncrawl/cc_chunk_embed.py at main · Eventual-Inc/daft-examples · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
# /// script
# description = "Chunk and embed Common Crawl text with spaCy and sentence-transformers"
# requires-python = ">=3.12, <3.13"
# dependencies = ["daft>=0.7.10", "torch", "sentence-transformers", "spacy", "pip", "python-dotenv"]
# ///

import spacy

import daft
from daft import DataType
from daft.functions import unnest

SpacyReturnType = DataType.list(
    DataType.struct(
        {
            "sent_id": DataType.int32(),
            "sent_start": DataType.int32(),
            "sent_end": DataType.int32(),
            "sent_text": DataType.string(),
            "sent_ents": DataType.list(DataType.string()),
        }
    )
)


@daft.cls()
class SpacyChunker:
    def __init__(self, model="en_core_web_sm"):
        self.nlp = spacy.load(model)

    @daft.method(return_dtype=SpacyReturnType)
    def chunk_text(self, text: str):
        doc = self.nlp(text)
        return [
            {
                "sent_id": i,
                "sent_start": sent.start,
                "sent_end": sent.end,
                "sent_text": sent.text,
                "sent_ents": [ent.text for ent in sent.ents],
            }
            for i, sent in enumerate(doc.sents)
        ]


if __name__ == "__main__":
    import daft
    from daft import col
    from daft.functions import decode, embed_text
    from daft.io import IOConfig, S3Config

    SOURCE_URI = "s3://commoncrawl/crawl-data/CC-MAIN-2025-33/segments/1754151579063.98/warc/CC-MAIN-20250815204238-20250815234238-00999.warc.gz"
    DEST_URI = ".data/common_crawl/chunk_embed"
    SPACY_MODEL = "en_core_web_sm"  # use "en_core_web_trf" for best accuracy (457.4 MB)
    EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
    NUM_FILES = 2

    IN_AWS = False
    IOCONFIG = IOConfig(s3=S3Config(anonymous=True, region_name="us-east-1"))

    # Read Preprocessed Text from Common Crawl WET
    df_warc = daft.datasets.common_crawl(
        "CC-MAIN-2025-33",
        content="text",
        num_files=NUM_FILES,
        in_aws=IN_AWS,
        io_config=IOCONFIG,  # Apply Creds
    ).limit(NUM_FILES)

    # Download the spaCy model
    spacy.cli.download(SPACY_MODEL)

    # Initialize the spaCy chunker
    spacy_chunk_text = SpacyChunker(model=SPACY_MODEL)

    # Chunk Text into sentences with spaCy
    df_prep = (
        df_warc.with_column(
            "warc_content",
            decode(col("warc_content"), "utf-8"),
        )
        .with_column(
            "spacy_results",
            spacy_chunk_text.chunk_text(text=col("warc_content")),
        )
        .explode("spacy_results")
    )

    # Embed Sentences
    df_embed = df_prep.with_column(
        "text_embeddings",
        embed_text(
            col("spacy_results")["sent_text"],
            model="sentence-transformers/all-MiniLM-L6-v2",
            provider="transformers",
        ),
    )

    df_embed.select(unnest(col("spacy_results")), col("text_embeddings")).show(format="fancy", max_width=20)