Upload 3 files

embedder.py

@@ -0,0 +1,241 @@
+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}")

llm_agent.py

@@ -0,0 +1,56 @@
+from openai import OpenAI
+from dotenv import load_dotenv
+import os
+from retriever import load_encoder, load_collection, encode_query, retrieve_docs, query_rerank, expand_with_neighbors, dedup_by_chapter_event
+from sentence_transformers import CrossEncoder
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+# load llm api key in .env
+load_dotenv()
+api_key = os.getenv("OPENAI_API_KEY")
+
+client = OpenAI(api_key=api_key)
+
+def build_rag_prompt(query, context):
+    prompt = f"""已知资料如下：
+{context}
+
+用户提问：{query}
+请参考所有已知资料, 并结合资料内容，简明、准确地回答问题。如果有多个符合的答案, 可以根据你是否确定而决定是否分别陈述这些答案.如果不能确定答案，请如实说明理由，不要凭空编造。"""
+    return prompt
+
+def llm_answer(query, expanded_results, model_name="gpt-4o"):
+    context = expanded_results[0][0] if expanded_results else ""
+    prompt = build_rag_prompt(query, context)
+    response = client.chat.completions.create(
+        model=model_name,
+        messages=[
+            {"role": "system", "content": "你是BangDream知识问答助手, 也就是邦学家. 只能基于提供的资料内容作答。"},
+            {"role": "user", "content": prompt}
+        ],
+        temperature=0.2,
+        max_tokens=512,
+    )
+    return response.choices[0].message.content.strip()
+
+if __name__ == "__main__":
+    collection = load_collection()
+    encoder = load_encoder()
+    reranker = CrossEncoder("BAAI/bge-reranker-large")
+
+    query_text = input("please enter your question：")
+    print("Thinking...\n...")
+    query_vec = encode_query(encoder, query_text)
+    results = retrieve_docs(collection, query_vec, top_k=50)
+    reranked = query_rerank(reranker, query_text, results, top_n=20)
+    deduped = dedup_by_chapter_event(reranked, max_per_group=1)
+    expanded_results = expand_with_neighbors(deduped[:5], collection)
+
+    answer = llm_answer(query_text, expanded_results)
+
+    print("\n=== Answer ===")
+    print(answer)
+    print("\n=== retrieved documents ===")
+    for idx, (context, score, meta) in enumerate(expanded_results, 1):
+        print(f"\n--- document {idx} (Score={score:.4f}) ---\n{context[:200]}...")
+        print(meta)

retriever.py

@@ -0,0 +1,174 @@
+import chromadb
+from sentence_transformers import SentenceTransformer, CrossEncoder
+
+CHROMA_DB_DIR = "./chroma_db"
+COLLECTION_NAME = "bangdream_dense"
+MODEL_NAME = "BAAI/bge-base-zh-v1.5"
+
+reranker = CrossEncoder("BAAI/bge-reranker-large")
+
+
+def load_collection(db_path=CHROMA_DB_DIR, collection_name=COLLECTION_NAME):
+    """Connect to Chroma persistent DB and load a collection."""
+    client = chromadb.PersistentClient(path=db_path)
+    collection = client.get_or_create_collection(collection_name)
+    return collection
+
+def load_encoder(model_name=MODEL_NAME):
+    """Load dense encoder model."""
+    return SentenceTransformer(model_name)
+
+def encode_query(encoder, query_text):
+    """Encode query text into normalized embedding."""
+    return encoder.encode_query([query_text], normalize_embeddings=True)
+
+def dedup_by_chapter_event(reranked_docs, max_per_group=1):
+    """de-duplicate when chapterTitle and eventName are identical"""
+    seen = {}
+    deduped = []
+    for doc, score, meta in reranked_docs:
+        key = (meta.get("chapterTitle", ""), meta.get("eventName", ""))
+        if key not in seen:
+            seen[key] = 1
+            deduped.append((doc, score, meta))
+        elif seen[key] < max_per_group:
+            seen[key] += 1
+            deduped.append((doc, score, meta))
+    return deduped
+
+def retrieve_docs(collection, query_vec, top_k=5):
+    """Retrieve documents from Chroma collection."""
+    results = collection.query(
+        query_embeddings=query_vec,
+        n_results=top_k,
+        include=["metadatas", "documents", "distances"],
+    )
+    return results
+
+def query_rerank(reranker, query, results, top_n=3):
+    """Use CrossEncoder to re-rank retrieved results."""
+    docs = results["documents"][0]
+    pairs = [(query, doc) for doc in docs]
+
+    # CrossEncoder
+    scores = reranker.predict(pairs)
+
+    # rerank
+    ranked = sorted(zip(docs, scores, results["metadatas"][0]), key=lambda x: x[1], reverse=True)
+
+    # get top_n
+    reranked_docs = ranked[:top_n]
+
+    """
+    # print result
+    print("=== After Rerank ===")
+    for i, (doc, score, meta) in enumerate(reranked_docs, 1):
+        print(f"Rank {i} | Score: {score:.4f}")
+        print(meta)
+        print(doc)
+        print("-" * 40)
+    """
+
+    return reranked_docs
+
+def pretty_print_results(results):
+    """Nicely print retrieved results."""
+    docs = results["documents"][0]
+    dists = results["distances"][0]
+    metas = results["metadatas"][0]
+    for idx, (doc, dist, meta) in enumerate(zip(docs, dists, metas)):
+        print(f"Rank {idx + 1} | Distance: {dist:.4f}")
+        print(meta)
+        print(doc)
+        print("-" * 40)
+
+# expend documents
+def get_all_chunks_in_chapter(collection, chapter_title, event_name=None, story_type=None):
+    filters = []
+    if chapter_title:
+        filters.append({"chapterTitle": chapter_title})
+    if story_type:
+        filters.append({"story_type": story_type})
+    if event_name:
+        filters.append({"eventName": event_name})
+    if len(filters) == 1:
+        filter_dict = filters[0]
+    elif len(filters) > 1:
+        filter_dict = {"$and": filters}
+    else:
+        filter_dict = {}
+    results = collection.get(where=filter_dict, include=["documents", "metadatas"])
+    chunk_list = []
+    for doc, meta in zip(results["documents"], results["metadatas"]):
+        chunk_list.append({
+            "text": doc,
+            **meta,
+        })
+    return chunk_list
+
+def find_adjacent_chunks(current_chunk, all_chunks):
+    start_idx = current_chunk['start_idx']
+    end_idx = current_chunk['end_idx']
+    prev_chunk, next_chunk = None, None
+    for chunk in all_chunks:
+        if chunk['end_idx'] == start_idx - 1:
+            prev_chunk = chunk
+        if chunk['start_idx'] == end_idx + 1:
+            next_chunk = chunk
+    return prev_chunk, next_chunk
+
+def safe_to_list(x):
+    if isinstance(x, str):
+        return x.split('\n') if '\n' in x else [x]
+    return list(x)
+
+def expand_with_neighbors(reranked_docs, collection):
+    expanded_results = []
+    for doc, score, meta in reranked_docs:
+        #print(meta)
+        chapter_title = meta.get("chapterTitle", "")
+        event_name = meta.get("eventName", "")
+        story_type = meta.get("story_type", None)
+        all_chunks = get_all_chunks_in_chapter(collection, chapter_title, event_name, story_type)
+        prev_chunk, next_chunk = find_adjacent_chunks(meta, all_chunks)
+        expanded_text = []
+        if prev_chunk:
+            #expanded_text += prev_chunk["text"]
+            expanded_text += safe_to_list(prev_chunk["text"])
+            #expanded_text.extend(prev_chunk["text"])
+        #expanded_text += doc
+        expanded_text += safe_to_list(doc)
+
+        #expanded_text.extend(doc if isinstance(doc, list) else [doc])
+        if next_chunk:
+            #expanded_text.extend(next_chunk["text"])
+            #expanded_text += next_chunk["text"]
+            expanded_text += safe_to_list(next_chunk["text"])
+
+        expanded_results.append((
+            "\n".join(expanded_text),
+            score,
+            {
+                **meta,
+                #"prev_chunk_id": prev_chunk["ids"][0] if prev_chunk else None,
+                #"next_chunk_id": next_chunk["ids"][0] if next_chunk else None,
+            }
+        ))
+    return expanded_results
+
+"""if __name__ == "__main__":
+    collection = load_collection()
+    encoder = load_encoder()
+
+    query_text = "乐奈喜欢什么?"
+    query_vec = encode_query(encoder, query_text)
+    results = retrieve_docs(collection, query_vec, top_k=50)
+    reranked = query_rerank(reranker, query_text, results, top_n=20)
+    deduped = dedup_by_chapter_event(reranked, max_per_group=1)
+    expanded_results = expand_with_neighbors(deduped[:5], collection)
+
+    for doc in expanded_results:
+        print("===")
+        print(doc)
+        print(doc[0])
+        print("===")"""