| 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") |
|
|
| |
| 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: |
| |
| chunked_docs.append({ |
| "text": row["QA_pair"], |
| "question": question, |
| "answer": ans, |
| "domain": domain, |
| "source_id": str(idx), |
| "chunk_num": "0" |
| }) |
| else: |
| |
| chunks = chunker.chunk(ans) |
| for chunk_num, chunk in enumerate(chunks): |
| chunked_docs.append({ |
| "text": f"Q: {question} A: {chunk.text}", |
| "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()) |
|
|
| |
| |
|
|
| |
| |
