RAG_Chatbot / data_preparation /chatbot_data_ingest.py
grazz7's picture
added chatbot codes
b22c324
Raw
History Blame Contribute Delete
2.07 kB
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())