原始文档: 79 个
分块后: 99 个

第一个分块示例:
硕士学位论 文
学位申请人姓名 闫腾
学位申请人学号 2200411007
专 业 名 称 机械工程
学 科 门 类 工学
学院（部、研究院） 应用技术学院
导 师 姓 名 李文贤
二〇二五年六月
分类号 学校代码 10590
UDC 密 级 公开
基于视-触觉融合感知的机器人
抓取滑动检测与力控研究
ധᎪն࿐...

from langchain_community.embeddings import HuggingFaceEmbeddings
from modelscope import snapshot_download
import os

# 创建本地模型目录
local_models_dir = "./Models"
os.makedirs(local_models_dir, exist_ok=True)

# 从 ModelScope 下载模型
model_id = "maidalun/bce-embedding-base_v1"
local_model_path = snapshot_download(model_id, cache_dir=local_models_dir)

# 初始化嵌入模型（启用 GPU）
embeddings = HuggingFaceEmbeddings(
    model_name=local_model_path,
    model_kwargs={"device": "cuda"},               # 使用 GPU 加速
    encode_kwargs={"normalize_embeddings": True}   # 归一化便于计算余弦相似度
)

# 测试单条文本向量化
text = "RAG是一种强大的AI技术"
vector = embeddings.embed_query(text)

print(f"文本: {text}")
print(f"向量维度: {len(vector)}")
print(f"向量前5个值: {[round(x, 4) for x in vector[:5]]}")

from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

text1 = "苹果是一种水果"
text2 = "香蕉是一种水果"
text3 = "苹果是一种好吃的水果"

# 获取向量
v1 = np.array(embeddings.embed_query(text1)).reshape(1, -1)
v2 = np.array(embeddings.embed_query(text2)).reshape(1, -1)
v3 = np.array(embeddings.embed_query(text3)).reshape(1, -1)

# 计算相似度
sim_1_2 = cosine_similarity(v1, v2)[0][0]
sim_1_3 = cosine_similarity(v1, v3)[0][0]

print(f"「{text1}」 vs 「{text2}」 相似度: {sim_1_2:.4f}")
print(f"「{text1}」 vs 「{text3}」 相似度: {sim_1_3:.4f}")

import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

def load_all_pdfs_from_directory(directory_path="./Dataset/PDF/"):
    """加载目录下所有PDF文件，并为每个分块添加唯一ID"""
    all_splits = []
    pdf_files = []
    
    if not os.path.exists(directory_path):
        print(f"❌ 目录不存在: {directory_path}")
        return all_splits, pdf_files
    
    # 获取所有 PDF 文件
    for filename in os.listdir(directory_path):
        if filename.lower().endswith('.pdf'):
            pdf_files.append(filename)
    
    if not pdf_files:
        print(f"⚠️  目录中没有找到PDF文件: {directory_path}")
        return all_splits, pdf_files
    
    print(f"📁 找到 {len(pdf_files)} 个PDF文件:")
    
    # 创建分块器（注意：此处 chunk_size=100 是示例，实际可调大）
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=100,
        chunk_overlap=20,
        length_function=len,
    )
    
    global_chunk_index = 0
    for filename in pdf_files:
        file_path = os.path.join(directory_path, filename)
        base_name = os.path.splitext(filename)[0]
        
        try:
            print(f"📖 正在加载: {filename}")
            loader = PyPDFLoader(file_path)
            docs = loader.load()
            splits = text_splitter.split_documents(docs)
            
            for i, split in enumerate(splits):
                page_num = split.metadata.get('page', 'unknown')
                split_id = f"{base_name}_p{page_num}_c{i}"
                split.metadata.update({
                    'id': split_id,
                    'source_file': filename
                })
                all_splits.append(split)
            
            print(f"✅ {filename}: {len(docs)}页 -> {len(splits)}个分块")
        except Exception as e:
            print(f"❌ 加载失败 {filename}: {str(e)}")
    
    print(f"🎉 总共加载 {len(all_splits)} 个文本分块（均已添加 id）")
    return all_splits, pdf_files

# 执行加载
splits, loaded_files = load_all_pdfs_from_directory("./Dataset/PDF/")

from langchain_chroma import Chroma

def smart_vectorstore_creation(documents, embeddings, persist_directory="./chroma_db", collection_name="langchain"):
    """智能向量数据库创建：自动清理旧数据，分批写入"""
    print(f"🔧 处理集合: {collection_name}")
    os.makedirs(persist_directory, exist_ok=True)
    
    try:
        # 尝试连接现有集合
        vectorstore = Chroma(
            collection_name=collection_name,
            persist_directory=persist_directory,
            embedding_function=embeddings
        )
        collection_info = vectorstore.get()
        if collection_info['ids']:
            print(f"🗑️ 清理集合中的 {len(collection_info['ids'])} 个文档")
            vectorstore.delete(ids=collection_info['ids'])
        else:
            print("✅ 集合为空，无需清理")
    except Exception:
        print(f"🆕 集合不存在，将创建新集合")
        vectorstore = None

    # 分批写入（Chroma 单次写入不宜过大）
    total_docs = len(documents)
    batch_size = 1000
    for i in range(0, total_docs, batch_size):
        batch = documents[i:i + batch_size]
        batch_num = i // batch_size + 1
        total_batches = (total_docs - 1) // batch_size + 1
        print(f"📦 添加批次 {batch_num}/{total_batches}: {len(batch)} 个文档")
        
        if vectorstore is None and i == 0:
            vectorstore = Chroma.from_documents(
                documents=batch,
                embedding=embeddings,
                persist_directory=persist_directory,
                collection_name=collection_name
            )
        else:
            vectorstore.add_documents(batch)
    
    print(f"✅ 完成！集合 '{collection_name}' 现有 {total_docs} 个文档")
    return vectorstore

# 执行存储
vectorstore = smart_vectorstore_creation(splits, embeddings, collection_name="robot_grasping_rag")

def custom_retriever(question: str) -> list:
    # 1. 扩大检索范围
    base_docs = vectorstore.similarity_search(question, k=10)
    
    # 2. 按来源过滤（例如只看“张静”的论文）
    filtered_docs = [
        doc for doc in base_docs
        if "张静" in doc.metadata.get('source_file', '')
    ]
    
    # 3. 重排序（此处简化为关键词匹配，实际可用 BCE-Reranker）
    def calculate_relevance_score(doc, query):
        return sum(1 for word in query if word in doc.page_content)
    
    scored_docs = [(doc, calculate_relevance_score(doc, question)) for doc in filtered_docs]
    scored_docs.sort(key=lambda x: x[1], reverse=True)
    
    # 4. 返回 top-5
    return [doc for doc, _ in scored_docs[:5]]

def evaluate_retrieval(retriever, test_cases):
    metrics = {"precision": [], "recall": []}
    for query, expected_ids in test_cases:
        retrieved = retriever.invoke(query)
        retrieved_ids = [d.metadata['id'] for d in retrieved]
        
        relevant = set(retrieved_ids) & set(expected_ids)
        precision = len(relevant) / len(retrieved_ids) if retrieved_ids else 0
        recall = len(relevant) / len(expected_ids) if expected_ids else 0
        
        metrics["precision"].append(precision)
        metrics["recall"].append(recall)
    
    return {
        "avg_precision": sum(metrics["precision"]) / len(metrics["precision"]),
        "avg_recall": sum(metrics["recall"]) / len(metrics["recall"])
    }

# 测试用例
test_cases = [
    ("什么是抓取检测？", [
        "基于视触感知协同的多指灵巧手抓取方法研究_张静_p164_c3266",
        "基于视触感知协同的机器人抓取技术研究_祝会龙_p39_c500",
        # ...其他相关 ID
    ]),
    ("什么是机械臂？", ["基于触觉图像序列的机器人抓取目标状态感知_韩筱_p25_c407"]),
]

results = evaluate_retrieval(retriever, test_cases)
print(f"平均精确率: {results['avg_precision']:.2%}")
print(f"平均召回率: {results['avg_recall']:.2%}")

def _build_prompt(self, question, context):
    cleaned_context = self._clean_context(context)
    
    return f"""<|im_start|>system
			你是一个专业的AI助手。请严格基于以下上下文信息回答问题：
			
			{cleaned_context}
			
			请遵循以下规则：
			1. 只使用上下文中的信息回答
			2. 如果上下文不包含相关信息，请回答"我不知道"
			3. 保持回答准确、简洁
			4. 不要编造信息<|im_end|>
			<|im_start|>user
			{question}<|im_end|>
			<|im_start|>assistant
			"""

🚀 开始测试问答系统
======================================================================

🎯 问题 1: 什么是机械臂？它有什么能力？
--------------------------------------------------
🔍 查询嵌入维度: 768
✅ 检索到 5 个相关文档
💡 答案: 机械臂是一种可以自动执行任务的机器人手臂，它具有精确度高、可重复性好、负载能力强、工作半径大、自由度多等能力。
📚 参考文档: 5 个

🎯 问题 2: 解释一下滑动检测的基本概念
--------------------------------------------------
🔍 查询嵌入维度: 768
✅ 检索到 5 个相关文档  
💡 答案: 滑动检测是一种用于检测物体是否发生滑动的算法。它通过比较帧间的标准差来判断物体是否发生滑动。如果帧间的标准差超过设定的阈值，那么就认为物体发生了滑动。滑动检测算法通常用于机械臂的抓取任务中，以防止物体在抓取过程中滑落。
📚 参考文档: 5 个

# 根据场景选择合适的链类型
chain_configs = {
    "default": {"chain_type": "stuff", "k": 5},
    "many_docs": {"chain_type": "map_reduce", "k": 20},
    "high_quality": {"chain_type": "refine", "k": 8},
    "most_relevant": {"chain_type": "map_rerank", "k": 10}
}

def select_chain_type(scenario, document_count):
    if document_count > 15:
        return "map_reduce"  # 文档太多，需要分治
    elif scenario == "precision_critical":
        return "map_rerank"  # 精度要求高
    elif scenario == "quality_first":
        return "refine"      # 质量优先
    else:
        return "stuff"       # 默认选择