chore: 添加vLLM相关配置与脚本 (#30)

更新.gitignore文件以忽略vLLM模型目录，新增docker-compose.yml、download_model.py和model_test.py文件，用于配置和测试vLLM服务。删除不再使用的magic_pdf_parser.py文件。

.gitignore

@@ -46,4 +46,5 @@ nltk_data/
 web/public/logo_secret.svg
 web/public/logo_old.svg
 web/public/logo.svg
-web/src/locales/zh.ts
+web/src/locales/zh.ts
+vllm/models

management/server/magic_pdf_parser.py

@@ -1,141 +0,0 @@
-import os
-from io import BytesIO
-
-from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
-from magic_pdf.data.dataset import PymuDocDataset
-from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
-from magic_pdf.config.enums import SupportedPdfParseMethod
-
-def process_pdf_with_magic(file_content, callback=None):
-    """
-    使用magic_pdf处理PDF文件
-    
-    Args:
-        file_content: PDF文件内容
-        callback: 回调函数，用于更新进度
-        
-    Returns:
-        解析后的内容列表
-    """
-    try:
-        from magic_pdf.processor import PDFProcessor
-        from magic_pdf.extractor import TextExtractor, ImageExtractor
-        
-        if callback:
-            callback(0.1, "初始化Magic PDF解析器")
-        
-        # 创建临时文件
-        temp_dir = os.path.join(os.getcwd(), "temp")
-        os.makedirs(temp_dir, exist_ok=True)
-        
-        temp_pdf_path = os.path.join(temp_dir, "temp.pdf")
-        with open(temp_pdf_path, "wb") as f:
-            f.write(file_content)
-        
-        if callback:
-            callback(0.2, "开始解析PDF")
-        
-        # 初始化处理器
-        processor = PDFProcessor(temp_pdf_path)
-        
-        if callback:
-            callback(0.3, "提取文本内容")
-        
-        # 提取文本
-        text_extractor = TextExtractor(processor)
-        text_content = text_extractor.extract()
-        
-        if callback:
-            callback(0.5, "提取图片内容")
-        
-        # 提取图片
-        image_extractor = ImageExtractor(processor)
-        images = image_extractor.extract()
-        
-        if callback:
-            callback(0.7, "组织解析结果")
-        
-        # 组织结果
-        content_list = []
-        
-        # 添加文本内容
-        for page_num, page_text in enumerate(text_content):
-            content_list.append({
-                "type": "text",
-                "page": page_num + 1,
-                "text": page_text
-            })
-        
-        # 添加图片内容
-        for i, img in enumerate(images):
-            content_list.append({
-                "type": "image",
-                "page": img.get("page", i + 1),
-                "image_path": img.get("path", ""),
-                "caption": img.get("caption", "")
-            })
-        
-        # 清理临时文件
-        try:
-            os.remove(temp_pdf_path)
-        except:
-            pass
-        
-        if callback:
-            callback(1.0, "PDF解析完成")
-        
-        return content_list
-        
-    except ImportError:
-        # 如果magic_pdf未安装，使用简单的文本提取
-        if callback:
-            callback(0.2, "Magic PDF未安装，使用备用方法")
-        
-        try:
-            import PyPDF2
-            
-            if callback:
-                callback(0.3, "使用PyPDF2提取文本")
-            
-            pdf_reader = PyPDF2.PdfReader(BytesIO(file_content))
-            content_list = []
-            
-            for i, page in enumerate(pdf_reader.pages):
-                if callback and i % 5 == 0:
-                    progress = 0.3 + (i / len(pdf_reader.pages)) * 0.6
-                    callback(progress, f"正在处理第 {i+1}/{len(pdf_reader.pages)} 页")
-                
-                text = page.extract_text()
-                if text:
-                    content_list.append({
-                        "type": "text",
-                        "page": i + 1,
-                        "text": text
-                    })
-            
-            if callback:
-                callback(0.9, "文本提取完成")
-            
-            return content_list
-            
-        except Exception as e:
-            if callback:
-                callback(0.5, f"PDF解析失败: {str(e)}")
-            
-            # 最简单的备用方案
-            return [{
-                "type": "text",
-                "page": 1,
-                "text": "无法解析PDF文件内容"
-            }]
-    
-    except Exception as e:
-        if callback:
-            callback(0.5, f"PDF解析失败: {str(e)}")
-        
-        # 出错时返回空列表
-        return [{
-            "type": "text",
-            "page": 1,
-            "text": f"解析失败: {str(e)}"
-        }]

vllm/docker-compose.yml

@@ -0,0 +1,50 @@
+services:
+  vllm-bge:
+    image: vllm/vllm-openai:latest
+    ipc: host
+    volumes:
+      - ./models/bge-m3:/models
+    command: [
+      "--model", "/models",
+      "--served-model-name", "bge-m3",
+      "--dtype", "float16",
+      "--gpu-memory-utilization", "0.9",
+    ]
+    ports:
+      - "8000:8000"
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              capabilities: [gpu]
+    networks:
+      - ragflow
+
+  vllm-deepseek:
+    image: vllm/vllm-openai:latest
+    ipc: host
+    volumes:
+      - ./models/DeepSeek-R1-1.5B:/models
+    command: [
+      "--model", "/models",
+      "--served-model-name", "deepseek-r1",
+      "--dtype", "float16",
+      "--tensor-parallel-size", "1",
+      "--max-model-len", "4096"
+    ]
+    ports:
+      - "8001:8000"
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              capabilities: [gpu]
+    networks:
+      - ragflow
+
+networks:
+  ragflow:
+    name: docker_ragflow 
+    driver: bridge

vllm/download_model.py

@@ -0,0 +1,34 @@
+import os
+from huggingface_hub import snapshot_download
+
+# 1. 设置镜像源（国内加速）
+# os.environ["HF_ENDPOINT"] = "https://mirrors.tuna.tsinghua.edu.cn/hugging-face/"
+
+# 2. 定义模型列表（名称 + 下载路径）
+models_to_download = [
+    {
+        "repo_id": "BAAI/bge-m3",  # Embedding 模型
+        "local_dir": os.path.expanduser("./models/bge-m3"),
+    },
+    {
+        "repo_id": "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",  # LLM 模型
+        "local_dir": os.path.expanduser("./models/DeepSeek-R1-1.5B"),
+    }
+]
+
+# 3. 遍历下载所有模型
+for model in models_to_download:
+    while True:  # 断点续传重试机制
+        try:
+            print(f"开始下载模型: {model['repo_id']} 到目录: {model['local_dir']}")
+            snapshot_download(
+                repo_id=model["repo_id"],
+                local_dir=model["local_dir"],
+                resume_download=True,  # 启用断点续传
+                force_download=False,  # 避免重复下载已有文件
+                token=None,            # 如需访问私有模型，替换为你的 token
+            )
+            print(f"模型 {model['repo_id']} 下载完成！")
+            break
+        except Exception as e:
+            print(f"下载失败: {e}, 重试中...")

vllm/model_test.py

@@ -0,0 +1,47 @@
+import requests
+from openai import OpenAI
+
+# 测试 embedding 模型 (vllm-bge)
+def test_embedding(model, text):
+    """测试嵌入模型"""
+    client = OpenAI(base_url="http://localhost:8000/v1", api_key="1")
+    
+    response = client.embeddings.create(
+        model=model,  # 使用支持嵌入的模型
+        input=text  # 需要嵌入的文本
+    )
+    
+    # 打印嵌入响应内容
+    print(f"Embedding response: {response}")
+    
+    if response and response.data:
+        print(f"Embedding: {response.data[0].embedding}")
+    else:
+        print("Failed to get embedding.")
+
+# 测试文本生成模型 (vllm-deepseek)
+def test_chat(model, prompt):
+    """测试文本生成模型"""
+    client = OpenAI(base_url="http://localhost:8001/v1", api_key="1")
+    
+    response = client.completions.create(
+        model=model,
+        prompt=prompt
+    )
+    
+    # 打印生成的文本
+    print(f"Chat response: {response.choices[0].text}")
+
+def main():
+    # 测试文本生成模型 deepseek-r1
+    prompt = "你好，今天的天气怎么样？"
+    print("Testing vllm-deepseek model for chat...")
+    test_chat("deepseek-r1", prompt)
+    
+    # 测试嵌入模型 bge-m3
+    embedding_text = "我喜欢编程，尤其是做AI模型。"
+    print("\nTesting vllm-bge model for embedding...")
+    test_embedding("bge-m3", embedding_text)
+
+if __name__ == "__main__":
+    main()