File size: 2,022 Bytes
1d9bd9b | 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 | #%%
import os
import json
import chromadb
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
CHUNKS_DIR = "/content/drive/MyDrive/updatedparentchunk/chunks/parent_child"
client = chromadb.PersistentClient(
path="/content/drive/MyDrive/chroma_bge_v2"
)
collection = client.get_or_create_collection(
"sci_judgments_bge_v2"
)
model = SentenceTransformer(
"BAAI/bge-small-en-v1.5"
)
def flatten_metadata(meta):
flat = {}
for k, v in meta.items():
if isinstance(v, (str, int, float, bool)):
flat[k] = v
elif isinstance(v, (list, dict)):
flat[k] = json.dumps(v)
else:
flat[k] = str(v)
return flat
files = [
f for f in os.listdir(CHUNKS_DIR)
if f.endswith(".json")
]
print("Files:", len(files))
all_ids = []
all_docs = []
all_meta = []
for filename in tqdm(files):
with open(
os.path.join(CHUNKS_DIR, filename),
encoding="utf-8"
) as f:
data = json.load(f)
parent = data["parent"]
all_ids.append(parent["chunk_id"])
all_docs.append(parent["text"])
all_meta.append(
flatten_metadata(parent["metadata"])
)
for child in data["children"]:
all_ids.append(child["chunk_id"])
all_docs.append(child["text"])
all_meta.append(
flatten_metadata(child["metadata"])
)
print("Total chunks:", len(all_ids))
BATCH_SIZE = 512
for start in tqdm(
range(0, len(all_ids), BATCH_SIZE),
desc="Embedding"
):
end = min(
start + BATCH_SIZE,
len(all_ids)
)
batch_docs = all_docs[start:end]
embeddings = model.encode(
batch_docs,
batch_size=128,
normalize_embeddings=True,
show_progress_bar=False
)
collection.add(
ids=all_ids[start:end],
embeddings=embeddings.tolist(),
documents=batch_docs,
metadatas=all_meta[start:end]
)
print("\nDONE")
print("Collection count:", collection.count()) |