本地知识库（3）：我在闲鱼买了台服务器，进行ai训练UE5

构建 Unreal Engine 5.6的源码的向量知识库

数据加载：从指定路径加载 UE 引擎和插件的 C++ 源码文件（.h 和.cpp）
文本分割：将源码按合理大小分割为代码块（避免因文件过大导致嵌入效率低）
向量生成：使用 GPU 加速的 SentenceTransformer 模型将代码块转换为向量（嵌入）
存储管理：将向量、代码块内容及元数据存入 ChromaDB 向量数据库

一开始输出时用的cpu，现在改的输出时用的gpu

import os
from tqdm import tqdm
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
import chromadb
import torch
import time
import requests
import logging
from typing import List
import numpy as np

# === 设置路径 ===
UE_ENGINE_PATH = "/home/veyvin/UnrealEngine/Engine/Source"
UE_PLUGIN_PATH = "/home/veyvin/UnrealEngine/Engine/Plugins"
CHROMA_DB_PATH = "./chroma_ue56_db_simple"

# === 代理设置 ===
PROXY_SETTINGS = {
    "http": "http://192.168.1.18:10808",
    "https": "http://192.168.1.18:10808",
}

# 禁用 ChromaDB 遥测和详细日志
os.environ["ANONYMIZED_TELEMETRY"] = "false"
logging.getLogger("chromadb").setLevel(logging.WARNING)

def setup_proxy_environment():
    """设置代理环境变量"""
    os.environ["HTTP_PROXY"] = PROXY_SETTINGS["http"]
    os.environ["HTTPS_PROXY"] = PROXY_SETTINGS["https"]
    os.environ["ALL_PROXY"] = PROXY_SETTINGS["https"]
    os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
    print("🔧 代理环境已设置")

def install_requirements():
    """安装必要的依赖包"""
    try:
        import sentence_transformers
    except ImportError:
        print("📦 安装 sentence-transformers...")
        os.system("pip install sentence-transformers")
        import sentence_transformers

class EfficientGPUEmbedder:
    def __init__(self, model_name='all-MiniLM-L6-v2'):
        setup_proxy_environment()
        install_requirements()
        from sentence_transformers import SentenceTransformer
        
        print("🚀 初始化 GPU 嵌入模型...")
        
        if not torch.cuda.is_available():
            raise RuntimeError("❌ CUDA 不可用")
        
        # 加载模型
        try:
            self.model = SentenceTransformer(model_name)
        except Exception as e:
            print(f"❌ 模型加载失败: {e}")
            raise
        
        # 移动到 GPU
        self.device = 'cuda'
        self.model = self.model.to(self.device)
        
        # 预热
        print("🔥 预热模型...")
        self.model.encode(["warmup"], device=self.device, batch_size=1, show_progress_bar=False)
        torch.cuda.synchronize()
        
        model_device = next(self.model.parameters()).device
        print(f"✅ 模型已加载到: {model_device}")
        
    def encode_batch(self, texts: List[str]) -> List[List[float]]:
        """批量编码文本"""
        if not texts:
            return []
        
        print(f"🔧 处理 {len(texts)} 个文本...")
        
        try:
            with torch.no_grad():
                embeddings = self.model.encode(
                    texts,
                    device=self.device,
                    convert_to_tensor=True,
                    show_progress_bar=True,
                    batch_size=8,
                    normalize_embeddings=True
                )
            
            result = embeddings.cpu().numpy().tolist()
            print(f"✅ 成功生成 {len(result)} 个嵌入向量")
            return result
            
        except Exception as e:
            print(f"❌ 嵌入生成失败: {e}")
            return []

def create_chroma_collection():
    """创建 ChromaDB 集合（简化版）"""
    print("💾 创建 ChromaDB 集合...")
    
    # 使用更简单的客户端配置
    client = chromadb.PersistentClient(
        path=CHROMA_DB_PATH,
        settings=chromadb.Settings(
            anonymized_telemetry=False,
            allow_reset=True
        )
    )
    
    # 清理旧集合
    try:
        client.delete_collection("ue_source_code")
        print("♻️ 已删除旧集合")
    except:
        pass
    
    # 创建新集合（不立即设置嵌入函数）
    collection = client.create_collection(
        name="ue_source_code",
        metadata={"hnsw:space": "cosine"}
    )
    
    print("✅ ChromaDB 集合创建完成")
    return client, collection

def load_all_source_files(engine_path, plugin_path):
    """加载UE源码和插件文件"""
    loaders = []
    for path in [engine_path, plugin_path]:
        print(f"📂 扫描路径: {path}")
        
        # C++ 源文件
        loaders.append(
            DirectoryLoader(
                path,
                glob="**/*.cpp",
                loader_cls=TextLoader,
                show_progress=True,
                silent_errors=True,
            )
        )
        
        # C++ 头文件
        loaders.append(
            DirectoryLoader(
                path,
                glob="**/*.h",
                loader_cls=TextLoader,
                show_progress=True,
                silent_errors=True,
            )
        )
    
    docs = []
    for loader in loaders:
        try:
            loader_docs = loader.load()
            docs.extend(loader_docs)
            print(f"✅ 从 {loader.path} 加载了 {len(loader_docs)} 个文件")
        except Exception as e:
            print(f"⚠️ 加载路径 {loader.path} 时出错: {e}")
            continue
    
    return docs

def check_gpu_compatibility():
    """检查 GPU 兼容性"""
    print("=" * 60)
    print("🔍 GPU 兼容性检查")
    print("=" * 60)
    
    print(f"1. PyTorch CUDA 可用: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"2. GPU 设备: {torch.cuda.get_device_name(0)}")
        
        # 测试 GPU 计算
        try:
            test_tensor = torch.randn(100, 100).cuda()
            result = test_tensor @ test_tensor.T
            print(f"3. GPU 计算测试: ✅ 通过")
            del test_tensor, result
            return True
        except Exception as e:
            print(f"❌ GPU 测试失败: {e}")
            return False
    else:
        print("❌ GPU 不可用")
        return False

def test_proxy_connection():
    """测试代理连接"""
    print("🔧 测试代理连接...")
    try:
        response = requests.get("https://huggingface.co", proxies=PROXY_SETTINGS, timeout=10)
        if response.status_code == 200:
            print("✅ 代理连接成功")
            return True
        else:
            print(f"⚠️ 代理连接异常，状态码: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ 代理连接失败: {e}")
        return False

if __name__ == "__main__":
    # 测试代理连接
    if not test_proxy_connection():
        use_proxy = input("❌ 代理连接失败，是否继续？(y/n): ")
        if use_proxy.lower() != 'y':
            exit(1)
    
    # 检查 GPU
    if not check_gpu_compatibility():
        print("❌ GPU 不可用，程序退出")
        exit(1)
    
    # 检查路径
    if not os.path.exists(UE_ENGINE_PATH):
        print(f"❌ 引擎路径不存在: {UE_ENGINE_PATH}")
        exit(1)
    if not os.path.exists(UE_PLUGIN_PATH):
        print(f"❌ 插件路径不存在: {UE_PLUGIN_PATH}")
        exit(1)
    
    print("\n" + "="*50)
    print("📚 第一阶段：加载文档")
    print("="*50)
    docs = load_all_source_files(UE_ENGINE_PATH, UE_PLUGIN_PATH)
    print(f"✅ 共加载 {len(docs)} 个源码文件")

    if len(docs) == 0:
        print("❌ 没有加载到任何文档，请检查路径")
        exit(1)

    print("🪚 正在分割文本片段...")
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,  # 更小的块大小
        chunk_overlap=80,
        separators=["\n\n", "\n", " ", ""]
    )
    chunks = splitter.split_documents(docs)
    print(f"✅ 分割为 {len(chunks)} 个代码块")

    # 先创建 ChromaDB 集合
    print("\n" + "="*50)
    print("💾 第二阶段：创建数据库")
    print("="*50)
    client, collection = create_chroma_collection()

    # 然后初始化嵌入器
    print("\n" + "="*50)
    print("🔧 第三阶段：初始化嵌入器")
    print("="*50)
    embedder = EfficientGPUEmbedder()

    # 手动处理嵌入和添加到集合
    print("\n" + "="*50)
    print("🚀 第四阶段：处理嵌入")
    print("="*50)
    
    batch_size = 16  # 小批次
    start_time = time.time()
    torch.cuda.empty_cache()
    
    successful_chunks = 0
    
    for i in tqdm(range(0, len(chunks), batch_size), desc="嵌入处理进度"):
        end_idx = min(i + batch_size, len(chunks))
        batch_chunks = chunks[i:end_idx]
        
        documents = [chunk.page_content for chunk in batch_chunks]
        metadatas = [chunk.metadata for chunk in batch_chunks]
        ids = [f"chunk_{i+j}" for j in range(len(batch_chunks))]
        
        # 生成嵌入
        try:
            embeddings = embedder.encode_batch(documents)
            
            if embeddings and len(embeddings) == len(documents):
                # 手动添加到集合
                collection.add(
                    embeddings=embeddings,
                    documents=documents,
                    metadatas=metadatas,
                    ids=ids
                )
                successful_chunks += len(batch_chunks)
            else:
                print(f"❌ 批次 {i} 嵌入生成不匹配")
                
        except Exception as e:
            print(f"❌ 批次 {i} 处理失败: {e}")
            continue
        
        # 定期清理内存
        if i % 100 == 0:
            torch.cuda.empty_cache()
    
    end_time = time.time()
    processing_time = end_time - start_time
    
    print("\n" + "="*50)
    print("🎉 知识库构建完成！")
    print("="*50)
    print(f"📊 数据库路径: {CHROMA_DB_PATH}")
    print(f"📚 总代码块: {len(chunks)}")
    print(f"✅ 成功导入: {successful_chunks}")
    print(f"⏱️ 处理时间: {processing_time:.2f} 秒")
    if processing_time > 0:
        print(f"⚡ 平均速度: {successful_chunks/processing_time:.2f} 文档/秒")

发现gpu利用率低时已经生成了一半了，所以没改，可以根据下面修改

🔍 主要原因分析

1. 批次大小过小

python

复制

下载

batch_size = 16  # 太小了！

P100有16GB显存，可以处理更大的批次
小批次无法充分利用并行计算能力

2. 模型规模限制

all-MiniLM-L6-v2 只有80MB，计算量不大
小模型无法占满P100的计算单元

3. 数据加载瓶颈

文件I/O和文本处理是CPU密集型
GPU在等待数据，造成闲置

🚀 优化方案

方案1：增大批次大小

python

复制

下载

# 修改这里
batch_size = 128  # 或者 256，根据显存调整

# 在encode_batch方法中也增大批次
def encode_batch(self, texts: List[str]) -> List[List[float]]:
    embeddings = self.model.encode(
        texts,
        device=self.device,
        convert_to_tensor=True,
        show_progress_bar=True,
        batch_size=64,  # 增大这个值
        normalize_embeddings=True
    )

方案2：使用更大模型

python

复制

下载

# 替换为更大的模型
class EfficientGPUEmbedder:
    def __init__(self, model_name='all-mpnet-base-v2'):  # 更大的模型
        # 或者使用以下模型：
        # 'multi-qa-mpnet-base-dot-v1'
        # 'all-distilroberta-v1'

方案3：流水线优化

python

复制

下载

# 预加载更多数据到内存
def process_in_larger_batches(self, chunks, batch_size=256):
    # 预先处理所有文档
    all_documents = [chunk.page_content for chunk in chunks]
    all_metadatas = [chunk.metadata for chunk in chunks]
    
    # 分更大的批次处理
    for i in tqdm(range(0, len(chunks), batch_size)):
        batch_texts = all_documents[i:i+batch_size]
        # ... 其余处理

方案4：启用TensorCore优化

python

复制

下载

# 确保使用混合精度
def encode_batch(self, texts: List[str]) -> List[List[float]]:
    with torch.no_grad():
        with torch.cuda.amp.autocast():  # 混合精度
            embeddings = self.model.encode(
                texts,
                device=self.device,
                batch_size=128,  # 大幅增加
                show_progress_bar=True
            )

📊 推荐的批次大小测试

python

复制

下载

# 测试不同批次大小
batch_sizes = [64, 128, 256, 512]
for bs in batch_sizes:
    torch.cuda.empty_cache()
    start_time = time.time()
    
    # 测试代码...
    print(f"批次大小 {bs}: 显存使用 {torch.cuda.memory_allocated()/1024**3:.1f}GB")

🎯 立即优化建议

首先尝试这个修改：

python

复制

下载

batch_size = 128  # 第214行附近

# 同时在encode_batch方法中：
batch_size=64  # 第78行附近

监控显存使用：

python

复制

下载

# 在处理循环中添加监控
if i % 10 == 0:
    gpu_memory = torch.cuda.memory_allocated() / 1024**3
    utilization = torch.cuda.utilization()
    print(f"GPU显存: {gpu_memory:.1f}GB, 利用率: {utilization}%")

P100有很强的计算能力，通过增大批次大小和使用更大模型，你应该能看到显存利用率提升到60-80%！

发现gpu利用率低时已经生成了一半了，所以没改，可以根据下面修改🔍 主要原因分析

1. 批次大小过小

2. 模型规模限制

3. 数据加载瓶颈

🚀 优化方案

方案1：增大批次大小

方案2：使用更大模型

方案3：流水线优化

方案4：启用TensorCore优化

📊 推荐的批次大小测试

🎯 立即优化建议

发现gpu利用率低时已经生成了一半了，所以没改，可以根据下面修改

🔍 主要原因分析