File size: 2,067 Bytes
b22c324 | 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 | from chonkie import SemanticChunker
from sentence_transformers import SentenceTransformer
from pathlib import Path
import pandas as pd
data_dir = Path("../datasets")
train_faq_df = pd.read_csv(data_dir / "processed" / "bank_faq" /"train_faq.csv")
# Initialize with an embedding model
chunker = SemanticChunker(
embedding_model="multi-qa-mpnet-base-dot-v1",
chunk_size=200
)
chunked_docs = []
for idx, row in train_faq_df.iterrows():
question = row["Question"]
ans = row["Answer"]
domain = row["Class"]
if len(ans.split()) < 200:
# short — keep as single doc, embed q+a together
chunked_docs.append({
"text": row["QA_pair"], # for embedding
"question": question,
"answer": ans,
"domain": domain,
"source_id": str(idx),
"chunk_num": "0"
})
else:
# long — chunk only the answer
chunks = chunker.chunk(ans)
for chunk_num, chunk in enumerate(chunks):
chunked_docs.append({
"text": f"Q: {question} A: {chunk.text}", # question anchors every chunk
"question": question,
"answer": chunk.text,
"domain": domain,
"source_id": str(idx),
"chunk_num": str(chunk_num)
})
print(f"Total chunks: {len(chunked_docs)}")
print(chunked_docs[0])
model = SentenceTransformer("multi-qa-mpnet-base-dot-v1")
texts = [doc["text"] for doc in chunked_docs]
embeddings = model.encode(texts, batch_size=64, show_progress_bar=True)
embed_faq_df = pd.DataFrame(chunked_docs)
embed_faq_df["embedding"] = embeddings.tolist()
embed_faq_df.to_csv(data_dir / "processed" / "bank_faq" / "embed_faq.csv", index=False)
print(embed_faq_df.shape)
print(embed_faq_df.head())
# for doc in chunked_docs:
# print(len(doc["text"]), "|", doc["text"][:80])
# cleaned_faq_df["token_count"] = cleaned_faq_df["QA_pair"].apply(lambda x: len(x.split()))
# print(cleaned_faq_df["token_count"].describe()) |