From dbf6fe0f3bbf0aea55e3d5460eb54d6514931509 Mon Sep 17 00:00:00 2001 From: zstar <65890619+zstar1003@users.noreply.github.com> Date: Wed, 16 Apr 2025 17:23:29 +0800 Subject: [PATCH] =?UTF-8?q?chore:=20=E6=B7=BB=E5=8A=A0vLLM=E7=9B=B8?= =?UTF-8?q?=E5=85=B3=E9=85=8D=E7=BD=AE=E4=B8=8E=E8=84=9A=E6=9C=AC=20(#30)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 更新.gitignore文件以忽略vLLM模型目录，新增docker-compose.yml、download_model.py和model_test.py文件，用于配置和测试vLLM服务。删除不再使用的magic_pdf_parser.py文件。 --- .gitignore | 3 +- management/server/magic_pdf_parser.py | 141 -------------------------- vllm/docker-compose.yml | 50 +++++++++ vllm/download_model.py | 34 +++++++ vllm/model_test.py | 47 +++++++++ 5 files changed, 133 insertions(+), 142 deletions(-) delete mode 100644 management/server/magic_pdf_parser.py create mode 100644 vllm/docker-compose.yml create mode 100644 vllm/download_model.py create mode 100644 vllm/model_test.py diff --git a/.gitignore b/.gitignore index 73e8704..db8418b 100644 --- a/.gitignore +++ b/.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 \ No newline at end of file +web/src/locales/zh.ts +vllm/models \ No newline at end of file diff --git a/management/server/magic_pdf_parser.py b/management/server/magic_pdf_parser.py deleted file mode 100644 index 14b82cc..0000000 --- a/management/server/magic_pdf_parser.py +++ /dev/null @@ -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)}" - }] \ No newline at end of file diff --git a/vllm/docker-compose.yml b/vllm/docker-compose.yml new file mode 100644 index 0000000..7cdd2ca --- /dev/null +++ b/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 \ No newline at end of file diff --git a/vllm/download_model.py b/vllm/download_model.py new file mode 100644 index 0000000..bd246ac --- /dev/null +++ b/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}, 重试中...") \ No newline at end of file diff --git a/vllm/model_test.py b/vllm/model_test.py new file mode 100644 index 0000000..780e8a6 --- /dev/null +++ b/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()