RagFlow 框架

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

# 前端缺少的依赖
@antv/g
@lexical/code
@lexical/rich-text
date-fns
katex
monaco-editor

# 配置环境变量
export UV_INDEX=https://mirrors.aliyun.com/pypi/simple
export HF_ENDPOINT=https://hf-mirror.com

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

#!/bin/bash

# 如果任何命令返回非 0 状态码，立即退出脚本
set -e

# ====================== #
# 函数：加载 .env 环境变量 #
# ====================== #
load_env_file() {
    # 获取当前脚本所在的目录
    # BASH_SOURCE 是一个数组，保存了当前 脚本/函数的源文件名，即使是 source 进来的也能正确拿到
    local script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    local env_file="$script_dir/.env"

    # 如果 .env 文件存在，就加载其中的环境变量
    if [ -f "$env_file" ]; then
        echo "Loading environment variables from: $env_file"
        # set -a = allexport 启用后，所有定义的变量都会自动被 export 到环境变量
        set -a               # 自动导出变量
        source "$env_file"   # 加载 .env 文件
        set +a
    else
        echo "Warning: .env file not found at: $env_file"
    fi
}

# 调用函数，加载环境变量
load_env_file

# 清空 http/https 代理（避免 Docker 守护进程遗留的代理影响程序）
export http_proxy=""; export https_proxy=""; export no_proxy=""
export HTTP_PROXY=""; export HTTPS_PROXY=""; export NO_PROXY=""

# 设置 Python 搜索路径为当前目录
export PYTHONPATH=$(pwd)

# 设置动态库搜索路径（避免找不到依赖库）
export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu/

# 使用 pkg-config 查找 jemalloc 库，并设置预加载路径
# 查询 jemalloc 库的安装路径（libdir，即库文件所在目录）
JEMALLOC_PATH=$(pkg-config --variable=libdir jemalloc)/libjemalloc.so

# 默认使用 python3 作为解释器
PY=python3

# 如果 WS（worker 数量）没有设置，或者小于 1，则默认设置为 1
if [[ -z "$WS" || $WS -lt 1 ]]; then
  WS=1
fi

# 每个进程最多重试 5 次
MAX_RETRIES=5

# 标志位：是否停止
STOP=false

# 用于记录子进程 PID 的数组
PIDS=()

# 设置 NLTK 数据目录（避免每次都联网下载）
export NLTK_DATA="./nltk_data"

# ============= #
# 函数：清理退出 #
# ============= #
cleanup() {
  echo "Termination signal received. Shutting down..."
  STOP=true
  # 遍历所有子进程 PID，逐个 kill
  for pid in "${PIDS[@]}"; do
    if kill -0 "$pid" 2>/dev/null; then
      echo "Killing process $pid"
      kill "$pid"
    fi
  done
  exit 0
}

# 捕获 SIGINT (Ctrl+C) 和 SIGTERM (系统停止)，调用 cleanup
trap cleanup SIGINT SIGTERM

# ================================================= #
# 函数：启动 task_executor.py，并带有重试逻辑        #
# 参数：task_id                                      #
# ================================================= #
task_exe(){
    # 将函数参数 $1 赋给局部变量 task_id
    local task_id=$1
    local retry_count=0
    while ! $STOP && [ $retry_count -lt $MAX_RETRIES ]; do
        echo "Starting task_executor.py for task $task_id (Attempt $((retry_count+1)))"
        # 启动任务执行器，并预加载 jemalloc 库，减少内存碎片
        LD_PRELOAD=$JEMALLOC_PATH $PY rag/svr/task_executor.py "$task_id"
        EXIT_CODE=$?
        if [ $EXIT_CODE -eq 0 ]; then
            echo "task_executor.py for task $task_id exited successfully."
            break
        else
            echo "task_executor.py for task $task_id failed with exit code $EXIT_CODE. Retrying..." >&2
            retry_count=$((retry_count + 1))
            sleep 2
        fi
    done

    # 如果超过最大重试次数，调用 cleanup 停止所有进程
    if [ $retry_count -ge $MAX_RETRIES ]; then
        echo "task_executor.py for task $task_id failed after $MAX_RETRIES attempts. Exiting..." >&2
        cleanup
    fi
}

# ================================================= #
# 函数：启动 ragflow_server.py，并带有重试逻辑       #
# ================================================= #
run_server(){
    local retry_count=0
    while ! $STOP && [ $retry_count -lt $MAX_RETRIES ]; do
        echo "Starting ragflow_server.py (Attempt $((retry_count+1)))"
        $PY api/ragflow_server.py
        EXIT_CODE=$?
        if [ $EXIT_CODE -eq 0 ]; then
            echo "ragflow_server.py exited successfully."
            break
        else
            echo "ragflow_server.py failed with exit code $EXIT_CODE. Retrying..." >&2
            retry_count=$((retry_count + 1))
            sleep 2
        fi
    done

    if [ $retry_count -ge $MAX_RETRIES ]; then
        echo "ragflow_server.py failed after $MAX_RETRIES attempts. Exiting..." >&2
        cleanup
    fi
}

# ======================= #
# 启动多个 task_executor  #
# ======================= #
for ((i=0;i<WS;i++))
do
  task_exe "$i" &      # 后台运行
  PIDS+=($!)           # 保存子进程 PID
done

# 启动主服务器
run_server &
PIDS+=($!)

# 等待所有子进程结束
wait

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

#!/usr/bin/env python3

# PEP 723 metadata
# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "huggingface-hub",
#   "nltk",
#   "argparse",
# ]
# ///

from huggingface_hub import snapshot_download
from typing import Union
import nltk
import os
import urllib.request
import argparse

def get_urls(use_china_mirrors=False) -> Union[str, list[str]]:
    if use_china_mirrors:
        return [
            "http://mirrors.tuna.tsinghua.edu.cn/ubuntu/pool/main/o/openssl/libssl1.1_1.1.1f-1ubuntu2_amd64.deb",
            "http://mirrors.tuna.tsinghua.edu.cn/ubuntu-ports/pool/main/o/openssl/libssl1.1_1.1.1f-1ubuntu2_arm64.deb",
            "https://repo.huaweicloud.com/repository/maven/org/apache/tika/tika-server-standard/3.0.0/tika-server-standard-3.0.0.jar",
            "https://repo.huaweicloud.com/repository/maven/org/apache/tika/tika-server-standard/3.0.0/tika-server-standard-3.0.0.jar.md5",
            "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken",
            ["https://registry.npmmirror.com/-/binary/chrome-for-testing/121.0.6167.85/linux64/chrome-linux64.zip", "chrome-linux64-121-0-6167-85"],
            ["https://registry.npmmirror.com/-/binary/chrome-for-testing/121.0.6167.85/linux64/chromedriver-linux64.zip", "chromedriver-linux64-121-0-6167-85"],
        ]
    else:
        return [
            "http://archive.ubuntu.com/ubuntu/pool/main/o/openssl/libssl1.1_1.1.1f-1ubuntu2_amd64.deb",
            "http://ports.ubuntu.com/pool/main/o/openssl/libssl1.1_1.1.1f-1ubuntu2_arm64.deb",
            "https://repo1.maven.org/maven2/org/apache/tika/tika-server-standard/3.0.0/tika-server-standard-3.0.0.jar",
            "https://repo1.maven.org/maven2/org/apache/tika/tika-server-standard/3.0.0/tika-server-standard-3.0.0.jar.md5",
            "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken",
            ["https://storage.googleapis.com/chrome-for-testing-public/121.0.6167.85/linux64/chrome-linux64.zip", "chrome-linux64-121-0-6167-85"],
            ["https://storage.googleapis.com/chrome-for-testing-public/121.0.6167.85/linux64/chromedriver-linux64.zip", "chromedriver-linux64-121-0-6167-85"],
        ]

repos = [
    "InfiniFlow/text_concat_xgb_v1.0",
    "InfiniFlow/deepdoc",
    "InfiniFlow/huqie",
    "BAAI/bge-large-zh-v1.5",
    "maidalun1020/bce-embedding-base_v1",
]

def download_model(repo_id):
    local_dir = os.path.abspath(os.path.join("huggingface.co", repo_id))
    os.makedirs(local_dir, exist_ok=True)
    snapshot_download(repo_id=repo_id, local_dir=local_dir)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Download dependencies with optional China mirror support')
    parser.add_argument('--china-mirrors', action='store_true', help='Use China-accessible mirrors for downloads')
    args = parser.parse_args()
    
    urls = get_urls(args.china_mirrors)
    
    for url in urls:
        download_url = url[0] if isinstance(url, list) else url
        filename = url[1] if isinstance(url, list) else url.split("/")[-1]
        print(f"Downloading {filename} from {download_url}...")
        if not os.path.exists(filename):
            urllib.request.urlretrieve(download_url, filename)

    local_dir = os.path.abspath('nltk_data')
    for data in ['wordnet', 'punkt', 'punkt_tab']:
        print(f"Downloading nltk {data}...")
        nltk.download(data, download_dir=local_dir)

    for repo_id in repos:
        print(f"Downloading huggingface repo {repo_id}...")
        download_model(repo_id)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

import tiktoken

# 加载编码
encoding = tiktoken.get_encoding("cl100k_base")

text = "Hello, world!"
tokens = encoding.encode(text)
print(tokens)  # 输出 token ID 列表

decoded = encoding.decode(tokens)
print(decoded)  # 输出原始文本 "Hello, world!"

1
2
3
4
5
6
7

import nltk

# 下载指定资源
nltk.download('wordnet')  # 下载 WordNet 词典

# 或者下载全部资源
nltk.download('all')

1

pip install huggingface_hub

1
2
3
4
5

from huggingface_hub import snapshot_download

# 下载整个模型仓库（默认到 ~/.cache/huggingface/hub/）
model_path = snapshot_download(repo_id="bert-base-uncased")
print(model_path)

1

snapshot_download(repo_id="bert-base-uncased", local_dir="./models/bert-base-uncased")

1

pip install transformers

1
2
3
4

from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("bert-base-uncased")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

1
2

git lfs install
git clone https://huggingface.co/bert-base-uncased

1
2
3
4
5
6
7
8

bert-base-uncased/
├── config.json            # 模型配置（层数、隐藏维度等）
├── pytorch_model.bin      # 模型权重（PyTorch）
├── tf_model.h5            # (可能有) TensorFlow 权重
├── tokenizer.json         # 分词器定义
├── tokenizer_config.json  # 分词器配置
├── vocab.txt              # 词表
└── special_tokens_map.json# 特殊 token 映射

1
2
3
4
5
6
7
8
9

from transformers import AutoModel, AutoTokenizer

model_path = "./models/bert-base-uncased"  # 本地路径
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path)

inputs = tokenizer("Hello Hugging Face!", return_tensors="pt")
outputs = model(**inputs)
print(outputs.last_hidden_state.shape)

1
2
3
4

from transformers import pipeline

classifier = pipeline("sentiment-analysis", model="./models/distilbert-base-uncased-finetuned-sst-2-english")
print(classifier("I love open-source models!"))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

from transformers import AutoTokenizer, AutoModelForCausalLM

model_path = "./models/meta-llama/Llama-2-7b-hf"

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")

prompt = "What is the capital of France?"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=50)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))