File size: 2,158 Bytes
dbabef2
0058d0d
dbabef2
 
0058d0d
 
dbabef2
70c8520
dbabef2
 
 
 
70c8520
dbabef2
70c8520
dbabef2
 
 
 
70c8520
dbabef2
 
 
70c8520
dbabef2
 
 
 
 
e755476
dbabef2
e755476
dbabef2
 
 
e755476
dbabef2
e755476
dbabef2
 
 
0058d0d
dbabef2
 
 
 
 
 
e755476
dbabef2
 
 
e755476
dbabef2
 
 
 
 
0058d0d
dbabef2
 
 
0058d0d
dbabef2
 
 
 
06407c6
dbabef2
0058d0d
dbabef2
 
 
e755476
dbabef2
e755476
dbabef2
 
 
e755476
dbabef2
0058d0d
dbabef2
 
 
 
 
 
 
 
 
 
e755476
dbabef2
 
 
 
e755476
dbabef2
 
0058d0d
dbabef2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
from __future__ import annotations

from db.embedder import LegalEmbedder
from db.vector_store import QdrantStore


class LegalIngestionPipeline:

    def __init__(
        self,
        collection_name: str = "legal_rag"
    ):

        self.embedder = LegalEmbedder()

        self.store = QdrantStore(
            collection_name=
                collection_name
        )

    # =====================================================
    # INGEST
    # =====================================================

    def ingest(
        self,
        chunks: list[dict],
        recreate_collection: bool = False
    ):

        if not chunks:

            print(
                "No chunks found."
            )

            return

        print(
            f"\nChunks: {len(chunks)}"
        )

        texts = [
            chunk[
                "enriched_text"
            ]
            for chunk in chunks
        ]

        print(
            "Generating embeddings..."
        )

        embeddings = (
            self.embedder.embed(
                texts
            )
        )

        vector_size = len(
            embeddings[0]
        )

        print(
            f"Vector Size: "
            f"{vector_size}"
        )

        if recreate_collection:

            self.store.recreate_collection(
                vector_size
            )

        else:

            self.store.create_collection(
                vector_size
            )

        points = []

        for idx, (
            chunk,
            vector
        ) in enumerate(
            zip(
                chunks,
                embeddings
            ),
            start=1
        ):

            points.append(
                {
                    "id":
                        idx,

                    "vector":
                        vector.tolist(),

                    "payload":
                        chunk
                }
            )

        print(
            f"Uploading "
            f"{len(points)} points..."
        )

        self.store.upsert_points(
            points
        )

        print(
            "\nIngestion Complete."
        )