ayrisdev/crawl4ai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: Update model_loader.py to create necessary folders in the home directory

Browse Source
This commit is contained in:
unclecode
2024-05-16 21:05:24 +08:00
parent 199c66114c
commit d19488a821
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
crawl4ai/model_loader.py
View File

@@ -1,11 +1,17 @@
from functools import lru_cache from functools import lru_cache
from .utils import get_home_folder
from pathlib import Path from pathlib import Path
import subprocess, os import subprocess, os
import shutil import shutil
from .config import MODEL_REPO_BRANCH from .config import MODEL_REPO_BRANCH
import argparse import argparse
def get_home_folder():
home_folder = os.path.join(Path.home(), ".crawl4ai")
os.makedirs(home_folder, exist_ok=True)
os.makedirs(f"{home_folder}/cache", exist_ok=True)
os.makedirs(f"{home_folder}/models", exist_ok=True)
return home_folder
@lru_cache() @lru_cache()
def load_bert_base_uncased(): def load_bert_base_uncased():
from transformers import BertTokenizer, BertModel, AutoTokenizer, AutoModel from transformers import BertTokenizer, BertModel, AutoTokenizer, AutoModel
@@ -89,6 +95,8 @@ def load_spacy_model():
def download_all_models(remove_existing=False): def download_all_models(remove_existing=False):
"""Download all models required for Crawl4AI.""" """Download all models required for Crawl4AI."""
print("[LOG] Welcome to the Crawl4AI Model Downloader!")
print("[LOG] This script will download all the models required for Crawl4AI.")
if remove_existing: if remove_existing:
print("[LOG] Removing existing models...") print("[LOG] Removing existing models...")
home_folder = get_home_folder() home_folder = get_home_folder()