Spaces:
Paused
Paused
# initializer.py | |
import json | |
import os | |
import requests | |
from huggingface_hub import snapshot_download | |
def download_json(url): | |
"""Download JSON content from the given URL.""" | |
response = requests.get(url) | |
response.raise_for_status() # Check if request was successful | |
return response.json() | |
def download_and_modify_json(url, local_filename, modifications): | |
"""Download JSON from URL and modify its contents based on modifications.""" | |
if os.path.exists(local_filename): | |
with open(local_filename, 'r', encoding='utf-8') as f: | |
data = json.load(f) | |
config_version = data.get('config_version', '0.0.0') | |
# If the version is older than desired, re-download | |
if config_version < '1.1.1': | |
data = download_json(url) | |
else: | |
data = download_json(url) | |
# Apply modifications | |
data.update(modifications) | |
# Save the modified JSON locally | |
with open(local_filename, 'w', encoding='utf-8') as f: | |
json.dump(data, f, ensure_ascii=False, indent=4) | |
def initialize_app(force=False): | |
"""Initialize models and configuration. | |
This function downloads the required models and modifies the configuration file. | |
It will only perform these actions if they have not been done before. | |
""" | |
# Define the path to the configuration file in the user's home directory. | |
home_dir = os.path.expanduser('~') | |
config_file_name = 'magic-pdf.json' | |
config_file = os.path.join(home_dir, config_file_name) | |
# If the config file exists, assume initialization is complete. | |
if os.path.exists(config_file): | |
print(f"Initialization already completed. Using existing configuration at {config_file}") | |
return | |
elif force: | |
print(f"Forced initialization.") | |
# Define patterns for model download | |
mineru_patterns = [ | |
"models/Layout/LayoutLMv3/*", | |
"models/Layout/YOLO/*", | |
"models/MFD/YOLO/*", | |
"models/MFR/unimernet_small_2501/*", | |
"models/TabRec/TableMaster/*", | |
"models/TabRec/StructEqTable/*", | |
] | |
model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', allow_patterns=mineru_patterns) | |
layoutreader_pattern = [ | |
"*.json", | |
"*.safetensors", | |
] | |
layoutreader_model_dir = snapshot_download('hantian/layoutreader', allow_patterns=layoutreader_pattern) | |
model_dir = os.path.join(model_dir, 'models') | |
print(f"Downloaded model_dir is: {model_dir}") | |
print(f"Downloaded layoutreader_model_dir is: {layoutreader_model_dir}") | |
# Download and modify JSON configuration | |
json_url = 'https://github.com/opendatalab/MinerU/raw/master/magic-pdf.template.json' | |
json_mods = { | |
'models-dir': model_dir, | |
'layoutreader-model-dir': layoutreader_model_dir, | |
} | |
download_and_modify_json(json_url, config_file, json_mods) | |
print(f"The configuration file has been configured successfully, the path is: {config_file}") | |
# This block will run if the module is executed directly. | |
if __name__ == '__main__': | |
initialize_app() | |