Upload 18 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/pdfs/Stream-Processing-with-Apache-Flink.pdf filter=lfs diff=lfs merge=lfs -text

config/rag_config.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+
+@dataclass
+class RAGConfig:
+    # Embedding 模型
+    embedding_model_name: str = "sentence-transformers/all-MiniLM-L6-v2"
+    normalize_embeddings: bool = True
+
+    # 检索参数
+    top_k: int = 5
+    similarity_threshold: float = 0.4
+
+    # LLM 模型
+    llm_model_name: str = "facebook/rag-token-base"
+    llm_max_length: int = 512
+    generation_kwargs: dict = None
+
+    # PDF 路径
+    pdf_dir: str = "data/pdfs"
+    vector_db_path: str = "data/embeddings/vector_store.pkl"
+
+    # Chunk 配置
+    chunk_size: int = 500
+    chunk_overlap: int = 100
+
+    def __post_init__(self):
+        if self.generation_kwargs is None:
+            self.generation_kwargs = {
+                "max_new_tokens": 200,
+                "temperature": 0.7,
+                "do_sample": True,
+            }

data/embeddings/vector_store.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0de8411ab4d21cfde6fcf8fc5db133064929163187c4c63203f0a842cf365df2
+size 3286005

data/pdfs/Stream-Processing-with-Apache-Flink.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f4804c75b17f898d45a811973ff188878d876a886dde28be43f0aaabed6bfc0
+size 10182829

prepare_embeddings.py

@@ -0,0 +1,30 @@
+import os
+import pickle
+from tqdm import tqdm
+from config.rag_config import RAGConfig
+from src.pdf_loader import load_pdfs, chunk_text
+from src.embedder import Embedder
+
+def prepare_embeddings():
+    config = RAGConfig()
+    embedder = Embedder(config)
+
+    print(f"📂 加载 PDF 文件...")
+    pdf_texts = load_pdfs(config.pdf_dir)
+
+    print("✂️ 切分文本...")
+    all_chunks = []
+    for text in pdf_texts:
+        all_chunks.extend(chunk_text(text, config.chunk_size, config.chunk_overlap))
+
+    print("🧠 生成 embeddings...")
+    embeddings = embedder.embed_texts(all_chunks)
+
+    os.makedirs(os.path.dirname(config.vector_db_path), exist_ok=True)
+    with open(config.vector_db_path, "wb") as f:
+        pickle.dump({"texts": all_chunks, "embeddings": embeddings}, f)
+
+    print(f"✅ 向量库已保存到 {config.vector_db_path}, 共 {len(all_chunks)} 段文本")
+
+if __name__ == "__main__":
+    prepare_embeddings()

requirements.txt

@@ -0,0 +1,35 @@
+# Transformers + RAG
+transformers>=4.43.0
+
+# Sentence Transformers 向量嵌入
+sentence-transformers>=2.2.2
+
+# Tokenizers
+tokenizers>=0.13.3
+
+# Hugging Face Hub
+huggingface_hub>=0.16.4
+datasets>=2.12.0
+
+# FAISS 向量检索
+faiss-cpu>=1.7.4
+
+# LangChain + 社区模块
+langchain>=0.2.19
+langchain-community>=0.1.16
+
+# PDF 解析
+pypdf>=3.15.0
+
+# FastAPI + REST API
+fastapi>=0.102.0
+uvicorn>=0.23.0
+
+torch
+transformers
+sentence-transformers
+scikit-learn
+gradio
+pypdf
+numpy
+tqdm

src/embedder.py

@@ -0,0 +1,14 @@
+from sentence_transformers import SentenceTransformer
+from sklearn.preprocessing import normalize
+from config.rag_config import RAGConfig
+
+class Embedder:
+    def __init__(self, config: RAGConfig):
+        self.model = SentenceTransformer(config.embedding_model_name)
+        self.normalize = config.normalize_embeddings
+
+    def embed_texts(self, texts):
+        embeddings = self.model.encode(texts, convert_to_numpy=True)
+        if self.normalize:
+            embeddings = normalize(embeddings)
+        return embeddings

src/pdf_loader.py

@@ -0,0 +1,23 @@
+import os
+from pypdf import PdfReader
+from config.rag_config import RAGConfig
+
+def load_pdfs(pdf_dir=None):
+    pdf_dir = pdf_dir or RAGConfig().pdf_dir
+    texts = []
+    for filename in os.listdir(pdf_dir):
+        if filename.endswith(".pdf"):
+            path = os.path.join(pdf_dir, filename)
+            reader = PdfReader(path)
+            text = "\n".join(page.extract_text() or "" for page in reader.pages)
+            texts.append(text)
+    return texts
+
+def chunk_text(text, chunk_size=500, overlap=100):
+    chunks = []
+    start = 0
+    while start < len(text):
+        end = start + chunk_size
+        chunks.append(text[start:end])
+        start += chunk_size - overlap
+    return chunks

src/rag_pipeline.py

@@ -0,0 +1,24 @@
+from transformers import RagTokenizer, RagSequenceForGeneration
+from config.rag_config import RAGConfig
+from src.embedder import Embedder
+from src.retriever import Retriever
+
+class RAGPipeline:
+    def __init__(self, config: RAGConfig, docs, doc_embeddings):
+        self.config = config
+        self.embedder = Embedder(config)
+        self.retriever = Retriever(doc_embeddings, docs, config)
+        self.tokenizer = RagTokenizer.from_pretrained(config.llm_model_name)
+        self.model = RagSequenceForGeneration.from_pretrained(config.llm_model_name)
+
+    def ask(self, query):
+        query_emb = self.embedder.embed_texts([query])[0]
+        retrieved = self.retriever.retrieve(query_emb)
+        context = "\n".join([r[0] for r in retrieved])
+        input_text = f"Question: {query}\nContext: {context}"
+        inputs = self.tokenizer(input_text, return_tensors="pt")
+        output = self.model.generate(
+            **inputs,
+            **self.config.generation_kwargs
+        )
+        return self.tokenizer.batch_decode(output, skip_special_tokens=True)[0], retrieved

src/retriever.py

@@ -0,0 +1,22 @@
+import numpy as np
+from sklearn.metrics.pairwise import cosine_similarity
+from config.rag_config import RAGConfig
+
+class Retriever:
+    def __init__(self, embeddings, texts, config: RAGConfig):
+        self.embeddings = embeddings
+        self.texts = texts
+        self.top_k = config.top_k
+        self.threshold = config.similarity_threshold
+
+    def retrieve(self, query_embedding):
+        scores = cosine_similarity([query_embedding], self.embeddings)[0]
+        # 阈值过滤
+        filtered = [(self.texts[i], float(scores[i]))
+                    for i in np.argsort(scores)[::-1]
+                    if scores[i] >= self.threshold]
+        results = filtered[:self.top_k]
+        if not results:
+            best_idx = int(np.argmax(scores))
+            results = [(self.texts[best_idx], float(scores[best_idx]))]
+        return results