from sentence_transformers import SentenceTransformer from sentence_transformers.sparse_encoder import SparseEncoder import torch import os import re from tqdm import tqdm import numpy as np import chromadb import json import time """ Currently only do dense encoding Sparse encoding related functions are placeholders """ DENSE_EMBEDDER_MODEL = "BAAI/bge-base-zh-v1.5" SPARSE_EMBEDDER_MODEL = "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-distill" class TextCleaner: def __init__(self, lowercase=False, remove_urls=True, normalize_space=True): self.lowercase = lowercase self.remove_urls = remove_urls self.normalize_space = normalize_space def clean(self, text: str) -> str: text = text.strip() if self.lowercase: text = text.lower() if self.remove_urls: text = re.sub(r"http\S+", "", text) if self.normalize_space: text = re.sub(r"\s+", " ", text) return text def join_chunk_text(text_chunk): # support chunk (list of sentences) processing if isinstance(text_chunk, list): return "\n".join(text_chunk) return text_chunk class DenseTextEncoder: """ output: numpy array """ def __init__(self, model_name, normalize=True, device=None): self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.model = SentenceTransformer(model_name) self.normalize = normalize self.cleaner = TextCleaner() def _prepare_texts(self, texts): """Support single string, list[str], or list[list[str]]""" if isinstance(texts, str): texts = [texts] elif isinstance(texts, list): if all(isinstance(t, str) for t in texts): texts = [join_chunk_text(t) if isinstance(t, list) else t for t in texts] elif all(isinstance(t, list) for t in texts): texts = [join_chunk_text(t) for t in texts] else: raise ValueError("Input list must contain only str or list[str].") else: raise ValueError("Input must be str or list.") cleaned = [self.cleaner.clean(t) for t in texts] return cleaned def encode_document(self, texts): cleaned = self._prepare_texts(texts) output = self.model.encode_document(cleaned, convert_to_numpy=True, normalize_embeddings=self.normalize) return output def encode_query(self, texts): cleaned = self._prepare_texts(texts) output = self.model.encode_query(cleaned, convert_to_numpy=True, normalize_embeddings=self.normalize) return output class SparseTextEncoder: """ output: torch tensor """ def __init__(self, model_name, device=None): self.device = device or "cpu" self.encoder = SparseEncoder(model_name, device=self.device) self.cleaner = TextCleaner() def _prepare_texts(self, texts): """Support single string, list[str], or list[list[str]]""" if isinstance(texts, str): texts = [texts] elif isinstance(texts, list): if all(isinstance(t, str) for t in texts): texts = [join_chunk_text(t) if isinstance(t, list) else t for t in texts] elif all(isinstance(t, list) for t in texts): texts = [join_chunk_text(t) for t in texts] else: raise ValueError("Input list must contain only str or list[str].") else: raise ValueError("Input must be str or list.") cleaned = [self.cleaner.clean(t) for t in texts] return cleaned def encode_document(self, texts): """Encode for corpus indexing""" cleaned = self._prepare_texts(texts) return self.encoder.encode_document(cleaned) def encode_query(self, texts): """Encode for query retrieval""" cleaned = self._prepare_texts(texts) return self.encoder.encode_query(cleaned) def read_input(source): if os.path.exists(source): with open(source, "r", encoding="utf-8") as f: lines = [line.strip() for line in f if line.strip()] return lines else: return [source] def encode_chunks_with_metadata(chunks, dense_encoder, sparse_encoder): """ :param chunks: [{'text': [...], 'chunk_id': ..., ...metadata}, ...] :param dense_encoder: dense encoder model :param sparse_encoder: sparse encoder model :return: { "chunk_id": "9b28e9938292486e9a61f2d1787bb828", "dense_embedding": np.array([...]), "sparse_embedding": torch.sparse.Tensor(...), "text": "友希那: ...\n莉莎: ...", "eventName": "连结思绪的未竟之歌", "chapterTitle": "序章: 古旧的磁带", "story_type": "event", # ...other metadata } """ text = [join_chunk_text(chunk["text"]) for chunk in chunks] dense_vecs = dense_encoder.encode_document(text) # placeholder, skip sparse encoding for now #sparse_vecs = sparse_encoder.encode_document(text) result = [] for i, chunk in enumerate(chunks): # placeholder, skip sparse encoding for now #sparse_i = sparse_vecs[i] #if isinstance(sparse_i, torch.Tensor) and sparse_i.is_sparse: # sparse_i = sparse_i.coalesce() result.append({ "chunk_id": chunk.get("chunk_id"), "dense_embedding": dense_vecs[i], # placeholder, skip sparse encoding for now "sparse_embedding": None, #"sparse_embedding": sparse_vecs[i], "text": text[i], "eventName": chunk.get("eventName"), "chapterTitle": chunk.get("chapterTitle"), "story_type": chunk.get("story_type"), "start_idx": chunk.get("start_idx"), "end_idx": chunk.get("end_idx"), }) return result # save dense embedding to chroma vector database def save_chunks_to_chroma(embedded_chunk, collection): ids = [] documents = [] embeddings = [] metadata = [] for entry in embedded_chunk: ids.append(entry["chunk_id"]) documents.append(entry["text"]) embeddings.append( entry["dense_embedding"].tolist() if isinstance(entry["dense_embedding"], np.ndarray) else entry[ "dense_embedding"]) # currently do not store sparse embedding to chroma meta = {k: v for k, v in entry.items() if k not in ["chunk_id", "dense_embedding","sparse_embedding", "text"]} metadata.append(meta) batch_size = 64 for i in range(0, len(ids), batch_size): collection.add( ids=ids[i:i + batch_size], documents=documents[i:i + batch_size], embeddings=embeddings[i:i + batch_size], metadatas=metadata[i:i + batch_size] ) print(f"saved {len(ids)} chunks to {collection.name}") def read_jsonl_in_batches(file_path, batch_size=64): batch = [] with open(file_path, 'r', encoding='utf8') as f: for line in f: if line.strip(): batch.append(json.loads(line)) if len(batch) == batch_size: yield batch batch = [] if batch: yield batch if __name__ == "__main__": chunk_files = [ "./chunks/band_chunks.jsonl", "./chunks/card_chunks.jsonl", "./chunks/event_chunks.jsonl", "./chunks/main_chunks.jsonl" ] dense_encoder = DenseTextEncoder(DENSE_EMBEDDER_MODEL) sparse_encoder = SparseTextEncoder(SPARSE_EMBEDDER_MODEL) # init databases chroma_client = chromadb.PersistentClient(path="./chroma_db") chroma_collection = chroma_client.get_or_create_collection("bangdream_dense") start_time = time.time() for file_path in chunk_files: with open(file_path, 'r', encoding='utf8') as f: total_lines = sum(1 for line in f if line.strip()) print(f"\nProcessing {file_path} ({total_lines} chunks)") pbar = tqdm(total=total_lines, desc=f"Encoding {os.path.basename(file_path)}", unit="chunk") for batch in read_jsonl_in_batches(file_path, batch_size=64): embedded = encode_chunks_with_metadata(batch, dense_encoder, sparse_encoder) save_chunks_to_chroma(embedded, chroma_collection) pbar.update(len(batch)) pbar.close() end_time = time.time() print(f"Total time used: {end_time - start_time}")