ragflow/download_deps.py
Jin Hai 996c94a8e7
Move clk100k_base tokenizer to docker image (#3411)
### What problem does this PR solve?

Move the tiktoken of cl100k_base into docker image

issue: #3338 

### Type of change

- [x] Refactoring

Signed-off-by: jinhai <haijin.chn@gmail.com>
Co-authored-by: Kevin Hu <kevinhu.sh@gmail.com>
2024-11-15 10:18:40 +08:00

48 lines
1.6 KiB
Python

#!/usr/bin/env python3
from os import rename
from huggingface_hub import snapshot_download
import nltk
import os
import urllib.request
urls = [
"http://archive.ubuntu.com/ubuntu/pool/main/o/openssl/libssl1.1_1.1.1f-1ubuntu2_amd64.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",
]
repos = [
"InfiniFlow/text_concat_xgb_v1.0",
"InfiniFlow/deepdoc",
"BAAI/bge-large-zh-v1.5",
"BAAI/bge-reranker-v2-m3",
"maidalun1020/bce-embedding-base_v1",
"maidalun1020/bce-reranker-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__":
for url in urls:
filename = url.split("/")[-1]
print(f"Downloading {url}...")
if not os.path.exists(filename):
urllib.request.urlretrieve(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)
rename("cl100k_base.tiktoken", "9b5ad71b2ce5302211f9c61530b329a4922fc6a4")