wehub-resource-sync
917eedffcf
Main / Python 3.11 - Docs (push) Has been cancelled
Main / Python 3.11 - Build (push) Has been cancelled
Main / Python 3.11 - Lint (push) Has been cancelled
Main / Python 3.11 - Style (push) Has been cancelled
Main / Python 3.11 - Test (push) Has been cancelled
Main / GPU CI (push) Has been cancelled
Main / Release (push) Has been cancelled
Main / Build and Push Docker Images (push) Has been cancelled
622 lines
25 KiB
Python
Executable File
622 lines
25 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Prepares OlmOCR checkpoints for deployment by:
|
|
1. Validating the model architecture
|
|
2. Copying model files to destination (disk or S3)
|
|
3. Downloading required tokenizer files from Hugging Face
|
|
|
|
Supports model souping (averaging weights of multiple checkpoints).
|
|
|
|
Usage:
|
|
python prepare_olmocr_checkpoint.py <source_path> <destination_path>
|
|
|
|
source_path: Path to checkpoint (local or S3)
|
|
destination_path: Where to save prepared checkpoint (local or S3)
|
|
|
|
For souping multiple checkpoints:
|
|
python prepare_olmocr_checkpoint.py <source1> <source2> ... <destination>
|
|
|
|
This will average the weights of all sources and prepare the souped checkpoint.
|
|
|
|
Examples:
|
|
# Single local to local
|
|
python prepare_olmocr_checkpoint.py /path/to/checkpoint /path/to/output
|
|
|
|
# Souping multiple S3 to S3
|
|
python prepare_olmocr_checkpoint.py s3://bucket/ckpt1 s3://bucket/ckpt2 s3://bucket/souped
|
|
|
|
# Mixed souping
|
|
python prepare_olmocr_checkpoint.py s3://bucket/ckpt1 /local/ckpt2 s3://bucket/souped
|
|
"""
|
|
|
|
import argparse
|
|
import concurrent.futures
|
|
import fnmatch
|
|
import json
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
from typing import Optional
|
|
|
|
import boto3
|
|
import requests
|
|
import torch
|
|
from botocore.exceptions import ClientError
|
|
from smart_open import smart_open
|
|
from tqdm import tqdm
|
|
from transformers import (
|
|
AutoConfig,
|
|
Qwen2_5_VLForConditionalGeneration,
|
|
Qwen2VLForConditionalGeneration,
|
|
)
|
|
|
|
try:
|
|
from safetensors.torch import load_file, save_file
|
|
except ImportError:
|
|
raise ImportError("Please install safetensors: pip install safetensors")
|
|
|
|
from olmocr.s3_utils import parse_s3_path
|
|
|
|
# Hugging Face model IDs for tokenizer files
|
|
HF_MODEL_IDS = {
|
|
"Qwen2VLForConditionalGeneration": "Qwen/Qwen2-VL-7B-Instruct",
|
|
"Qwen2_5_VLForConditionalGeneration": "Qwen/Qwen2.5-VL-7B-Instruct",
|
|
"Qwen3VLForConditionalGeneration": "Qwen/Qwen3-VL-8B-Instruct",
|
|
}
|
|
|
|
# Required tokenizer files to download from Hugging Face
|
|
TOKENIZER_FILES = ["chat_template.json", "merges.txt", "preprocessor_config.json", "tokenizer.json", "tokenizer_config.json", "vocab.json"]
|
|
|
|
# Supported model architectures
|
|
SUPPORTED_ARCHITECTURES = ["Qwen2VLForConditionalGeneration", "Qwen2_5_VLForConditionalGeneration", "Qwen3VLForConditionalGeneration"]
|
|
|
|
# Map architectures to corresponding model classes
|
|
MODEL_CLASS_MAP = {
|
|
"Qwen2VLForConditionalGeneration": Qwen2VLForConditionalGeneration,
|
|
"Qwen2_5_VLForConditionalGeneration": Qwen2_5_VLForConditionalGeneration,
|
|
}
|
|
|
|
# Files to exclude from copying (training-related files)
|
|
# Supports exact matches and glob patterns
|
|
EXCLUDED_FILES = {"optimizer.pt", "scheduler.pt", "rng_state.pth", "trainer_state.json", "training_args.bin", "*.pt", "*.pth"}
|
|
|
|
s3_client = boto3.client("s3")
|
|
|
|
|
|
def should_exclude_file(filename: str) -> bool:
|
|
"""Check if a file should be excluded based on EXCLUDED_FILES patterns."""
|
|
for pattern in EXCLUDED_FILES:
|
|
if fnmatch.fnmatch(filename, pattern):
|
|
return True
|
|
return False
|
|
|
|
|
|
def is_s3_path(path: str) -> bool:
|
|
"""Check if a path is an S3 path."""
|
|
return path.startswith("s3://")
|
|
|
|
|
|
def join_path(base: str, *parts: str) -> str:
|
|
"""Join paths for local and S3-style URIs."""
|
|
if not parts:
|
|
return base
|
|
if is_s3_path(base):
|
|
cleaned = [base.rstrip("/")]
|
|
for part in parts:
|
|
cleaned.append(part.strip("/"))
|
|
return "/".join(segment for segment in cleaned if segment)
|
|
return os.path.join(base, *parts)
|
|
|
|
|
|
def load_json_if_exists(path: str) -> Optional[dict]:
|
|
"""Load JSON from a path if it exists, otherwise return None."""
|
|
try:
|
|
with smart_open(path, "r") as handle:
|
|
return json.load(handle)
|
|
except FileNotFoundError:
|
|
return None
|
|
except ClientError as exc:
|
|
error_code = exc.response.get("Error", {}).get("Code")
|
|
if error_code in {"NoSuchKey", "404"}:
|
|
return None
|
|
raise
|
|
except OSError as exc:
|
|
# Handle S3 NoSuchKey errors that come through as OSError
|
|
exc_str = str(exc)
|
|
if "No such file" in exc_str:
|
|
return None
|
|
if "NoSuchKey" in exc_str:
|
|
return None
|
|
if "The specified key does not exist" in exc_str:
|
|
return None
|
|
raise
|
|
except Exception as exc:
|
|
# Catch any other S3-related errors for missing files
|
|
exc_str = str(exc)
|
|
if "NoSuchKey" in exc_str or "specified key does not exist" in exc_str.lower():
|
|
return None
|
|
raise
|
|
|
|
|
|
def load_adapter_config(source_path: str) -> Optional[dict]:
|
|
"""Return the LoRA adapter configuration if present for the given source."""
|
|
adapter_config_path = join_path(source_path, "adapter_config.json")
|
|
return load_json_if_exists(adapter_config_path)
|
|
|
|
|
|
def download_file_from_hf(filename: str, destination_dir: str, hf_base_url: str) -> None:
|
|
"""Download a file from Hugging Face model repository."""
|
|
url = f"{hf_base_url}/{filename}"
|
|
local_path = os.path.join(destination_dir, filename)
|
|
|
|
print(f"Downloading {filename} from Hugging Face...")
|
|
response = requests.get(url, stream=True)
|
|
response.raise_for_status()
|
|
|
|
with open(local_path, "wb") as f:
|
|
for chunk in response.iter_content(chunk_size=8192):
|
|
f.write(chunk)
|
|
|
|
print(f"Downloaded {filename}")
|
|
|
|
|
|
def detect_checkpoint_architecture(config_path: str) -> str:
|
|
"""Detect and validate the checkpoint architecture."""
|
|
print(f"Detecting checkpoint architecture from {config_path}...")
|
|
|
|
with smart_open(config_path, "r") as f:
|
|
config_data = json.load(f)
|
|
|
|
architectures = config_data.get("architectures", [])
|
|
|
|
# Find the supported architecture
|
|
detected_architecture = None
|
|
for arch in architectures:
|
|
if arch in SUPPORTED_ARCHITECTURES:
|
|
detected_architecture = arch
|
|
break
|
|
|
|
if not detected_architecture:
|
|
# Try to detect from model name
|
|
model_name = config_data.get("name_or_path", "")
|
|
if "Qwen2.5-VL" in model_name:
|
|
detected_architecture = "Qwen2_5_VLForConditionalGeneration"
|
|
elif "Qwen2-VL" in model_name:
|
|
detected_architecture = "Qwen2VLForConditionalGeneration"
|
|
elif "Qwen3-VL" in model_name:
|
|
detected_architecture = "Qwen3VLForConditionalGeneration"
|
|
else:
|
|
raise ValueError(f"No supported architecture found. Expected one of {SUPPORTED_ARCHITECTURES} " f"but found: {architectures}")
|
|
|
|
print(f"✓ Detected architecture: {detected_architecture}")
|
|
return detected_architecture
|
|
|
|
|
|
def copy_local_to_local(source_dir: str, dest_dir: str) -> None:
|
|
"""Copy files from local directory to local directory."""
|
|
os.makedirs(dest_dir, exist_ok=True)
|
|
|
|
# Get list of files to copy
|
|
files_to_copy = []
|
|
for root, _, files in os.walk(source_dir):
|
|
for file in files:
|
|
if should_exclude_file(file):
|
|
print(f"Skipping excluded file: {file}")
|
|
continue
|
|
src_path = os.path.join(root, file)
|
|
rel_path = os.path.relpath(src_path, source_dir)
|
|
files_to_copy.append((src_path, os.path.join(dest_dir, rel_path)))
|
|
|
|
print(f"Copying {len(files_to_copy)} files from {source_dir} to {dest_dir}...")
|
|
|
|
for src_path, dst_path in tqdm(files_to_copy, desc="Copying files"):
|
|
os.makedirs(os.path.dirname(dst_path), exist_ok=True)
|
|
shutil.copy2(src_path, dst_path)
|
|
|
|
|
|
def download_file_from_s3(bucket: str, key: str, local_path: str) -> None:
|
|
"""Download a single file from S3."""
|
|
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
|
s3_client.download_file(bucket, key, local_path)
|
|
|
|
|
|
def upload_file_to_s3(local_path: str, bucket: str, key: str) -> None:
|
|
"""Upload a single file to S3."""
|
|
s3_client.upload_file(local_path, bucket, key)
|
|
|
|
|
|
def copy_s3_to_local(source_bucket: str, source_prefix: str, dest_dir: str) -> None:
|
|
"""Copy files from S3 to local directory."""
|
|
os.makedirs(dest_dir, exist_ok=True)
|
|
|
|
# List all objects in source
|
|
paginator = s3_client.get_paginator("list_objects_v2")
|
|
pages = paginator.paginate(Bucket=source_bucket, Prefix=source_prefix)
|
|
|
|
download_tasks = []
|
|
for page in pages:
|
|
for obj in page.get("Contents", []):
|
|
key = obj["Key"]
|
|
if key.endswith("/"):
|
|
continue
|
|
|
|
filename = os.path.basename(key)
|
|
if should_exclude_file(filename):
|
|
print(f"Skipping excluded file: {filename}")
|
|
continue
|
|
|
|
rel_path = os.path.relpath(key, source_prefix)
|
|
local_path = os.path.join(dest_dir, rel_path)
|
|
download_tasks.append((source_bucket, key, local_path))
|
|
|
|
print(f"Downloading {len(download_tasks)} files from s3://{source_bucket}/{source_prefix} to {dest_dir}...")
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
|
|
futures = [executor.submit(download_file_from_s3, bucket, key, local_path) for bucket, key, local_path in download_tasks]
|
|
|
|
for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Downloading"):
|
|
future.result()
|
|
|
|
|
|
def copy_local_to_s3(source_dir: str, dest_bucket: str, dest_prefix: str) -> None:
|
|
"""Copy files from local directory to S3."""
|
|
# Get list of files to upload
|
|
upload_tasks = []
|
|
for root, _, files in os.walk(source_dir):
|
|
for file in files:
|
|
if should_exclude_file(file):
|
|
print(f"Skipping excluded file: {file}")
|
|
continue
|
|
local_path = os.path.join(root, file)
|
|
rel_path = os.path.relpath(local_path, source_dir)
|
|
s3_key = os.path.join(dest_prefix, rel_path)
|
|
upload_tasks.append((local_path, dest_bucket, s3_key))
|
|
|
|
print(f"Uploading {len(upload_tasks)} files from {source_dir} to s3://{dest_bucket}/{dest_prefix}...")
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
|
|
futures = [executor.submit(upload_file_to_s3, local_path, bucket, key) for local_path, bucket, key in upload_tasks]
|
|
|
|
for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Uploading"):
|
|
future.result()
|
|
|
|
|
|
def copy_s3_to_s3(source_bucket: str, source_prefix: str, dest_bucket: str, dest_prefix: str) -> None:
|
|
"""Copy files from S3 to S3."""
|
|
# List all objects in source
|
|
paginator = s3_client.get_paginator("list_objects_v2")
|
|
pages = paginator.paginate(Bucket=source_bucket, Prefix=source_prefix)
|
|
|
|
copy_tasks = []
|
|
for page in pages:
|
|
for obj in page.get("Contents", []):
|
|
key = obj["Key"]
|
|
if key.endswith("/"):
|
|
continue
|
|
|
|
filename = os.path.basename(key)
|
|
if should_exclude_file(filename):
|
|
print(f"Skipping excluded file: {filename}")
|
|
continue
|
|
|
|
rel_path = os.path.relpath(key, source_prefix)
|
|
dest_key = os.path.join(dest_prefix, rel_path)
|
|
copy_source = {"Bucket": source_bucket, "Key": key}
|
|
copy_tasks.append((copy_source, dest_bucket, dest_key))
|
|
|
|
print(f"Copying {len(copy_tasks)} files from s3://{source_bucket}/{source_prefix} to s3://{dest_bucket}/{dest_prefix}...")
|
|
|
|
for copy_source, bucket, key in tqdm(copy_tasks, desc="Copying"):
|
|
s3_client.copy_object(CopySource=copy_source, Bucket=bucket, Key=key)
|
|
|
|
|
|
def get_weight_files(dir_path: str) -> list[str]:
|
|
"""Get list of weight files (full paths) in the directory."""
|
|
weight_files = []
|
|
for root, _, files in os.walk(dir_path):
|
|
for file in files:
|
|
full_path = os.path.join(root, file)
|
|
if (file.startswith("pytorch_model") and file.endswith(".bin")) or file.endswith(".safetensors"):
|
|
weight_files.append(full_path)
|
|
return weight_files
|
|
|
|
|
|
def merge_lora_adapter_checkpoint(adapter_dir: str, base_model_name: str, output_dir: str) -> str:
|
|
"""Merge a LoRA adapter into its base model and save the merged weights.
|
|
|
|
Returns:
|
|
The detected architecture string of the merged model.
|
|
"""
|
|
try:
|
|
from peft import PeftModel
|
|
except ImportError as exc: # pragma: no cover - optional dependency guard
|
|
raise ImportError("Merging LoRA adapters requires the `peft` package. Install it with `pip install peft`.") from exc
|
|
|
|
print(f"Merging LoRA adapter from {adapter_dir} with base model '{base_model_name}'...")
|
|
base_config = AutoConfig.from_pretrained(base_model_name, trust_remote_code=True)
|
|
|
|
architecture = None
|
|
for arch in base_config.architectures or []:
|
|
if arch in MODEL_CLASS_MAP:
|
|
architecture = arch
|
|
break
|
|
|
|
if architecture is None:
|
|
raise ValueError(
|
|
f"Base model '{base_model_name}' uses an unsupported architecture: {base_config.architectures}. "
|
|
f"Supported architectures: {SUPPORTED_ARCHITECTURES}"
|
|
)
|
|
|
|
model_class = MODEL_CLASS_MAP[architecture]
|
|
base_model = model_class.from_pretrained(
|
|
base_model_name,
|
|
trust_remote_code=True,
|
|
torch_dtype="auto",
|
|
)
|
|
|
|
lora_model = PeftModel.from_pretrained(base_model, adapter_dir, is_trainable=False)
|
|
merged_model = lora_model.merge_and_unload()
|
|
merged_model = merged_model.to("cpu")
|
|
|
|
if hasattr(merged_model, "config"):
|
|
merged_model.config._name_or_path = base_model_name
|
|
merged_model.config.base_model_name_or_path = base_model_name
|
|
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
merged_model.save_pretrained(output_dir)
|
|
print(f"✓ Saved merged model to {output_dir}")
|
|
|
|
# Explicit cleanup
|
|
del merged_model
|
|
del lora_model
|
|
del base_model
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
return architecture
|
|
|
|
|
|
def prepare_checkpoints(sources: list[str], dest_path: str) -> None:
|
|
"""Prepare OlmOCR checkpoint(s) for deployment, with support for souping."""
|
|
print(f"Preparing {'souped ' if len(sources) > 1 else ''}checkpoint from {len(sources)} source(s) to {dest_path}")
|
|
|
|
sources = [source.rstrip("/") for source in sources]
|
|
dest_path = dest_path.rstrip("/")
|
|
|
|
source_infos = []
|
|
for source in sources:
|
|
adapter_config = load_adapter_config(source)
|
|
if adapter_config is not None:
|
|
source_infos.append({"path": source, "is_lora": True, "adapter_config": adapter_config})
|
|
else:
|
|
config_path = join_path(source, "config.json")
|
|
arch = detect_checkpoint_architecture(config_path)
|
|
source_infos.append({"path": source, "is_lora": False, "architecture": arch})
|
|
|
|
num_lora_sources = sum(1 for info in source_infos if info["is_lora"])
|
|
final_architecture: Optional[str] = None
|
|
|
|
if num_lora_sources > 0:
|
|
if len(source_infos) > 1:
|
|
raise ValueError("LoRA adapter checkpoints can only be processed individually, not during souping.")
|
|
|
|
source_info = source_infos[0]
|
|
source_path = source_info["path"]
|
|
adapter_config = source_info["adapter_config"]
|
|
base_model_name = adapter_config.get("base_model_name_or_path")
|
|
if not base_model_name:
|
|
raise ValueError("adapter_config.json is missing 'base_model_name_or_path'; cannot merge LoRA adapter.")
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
adapter_local_dir = os.path.join(temp_dir, "adapter")
|
|
print("\nDownloading LoRA adapter locally for merging...")
|
|
if is_s3_path(source_path):
|
|
bucket, prefix = parse_s3_path(source_path)
|
|
copy_s3_to_local(bucket, prefix, adapter_local_dir)
|
|
else:
|
|
copy_local_to_local(source_path, adapter_local_dir)
|
|
|
|
merged_dir = os.path.join(temp_dir, "merged")
|
|
final_architecture = merge_lora_adapter_checkpoint(adapter_local_dir, base_model_name, merged_dir)
|
|
|
|
print("\nCopying merged model files to destination...")
|
|
if is_s3_path(dest_path):
|
|
dest_bucket, dest_prefix = parse_s3_path(dest_path)
|
|
copy_local_to_s3(merged_dir, dest_bucket, dest_prefix)
|
|
else:
|
|
copy_local_to_local(merged_dir, dest_path)
|
|
else:
|
|
architectures = [info["architecture"] for info in source_infos]
|
|
if len(set(architectures)) > 1:
|
|
raise ValueError("All sources must have the same architecture")
|
|
|
|
final_architecture = architectures[0]
|
|
|
|
if len(sources) == 1:
|
|
source_path = sources[0]
|
|
print("\nCopying model files...")
|
|
if is_s3_path(source_path) and is_s3_path(dest_path):
|
|
source_bucket, source_prefix = parse_s3_path(source_path)
|
|
dest_bucket, dest_prefix = parse_s3_path(dest_path)
|
|
copy_s3_to_s3(source_bucket, source_prefix, dest_bucket, dest_prefix)
|
|
elif is_s3_path(source_path) and not is_s3_path(dest_path):
|
|
source_bucket, source_prefix = parse_s3_path(source_path)
|
|
copy_s3_to_local(source_bucket, source_prefix, dest_path)
|
|
elif not is_s3_path(source_path) and is_s3_path(dest_path):
|
|
dest_bucket, dest_prefix = parse_s3_path(dest_path)
|
|
copy_local_to_s3(source_path, dest_bucket, dest_prefix)
|
|
else:
|
|
copy_local_to_local(source_path, dest_path)
|
|
else:
|
|
# Souping multiple checkpoints
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
# Download all sources to local temp dirs
|
|
source_temps = []
|
|
for i, source in enumerate(sources):
|
|
source_temp = os.path.join(temp_dir, f"source_{i}")
|
|
if is_s3_path(source):
|
|
bucket, prefix = parse_s3_path(source)
|
|
copy_s3_to_local(bucket, prefix, source_temp)
|
|
else:
|
|
copy_local_to_local(source, source_temp)
|
|
source_temps.append(source_temp)
|
|
|
|
first_source = source_temps[0]
|
|
|
|
# Get weight files
|
|
weight_full_paths = get_weight_files(first_source)
|
|
weight_rel_paths = [os.path.relpath(p, first_source) for p in weight_full_paths]
|
|
|
|
# Verify others have same weight files
|
|
for i in range(1, len(sources)):
|
|
other_dir = source_temps[i]
|
|
other_weights = [os.path.relpath(p, other_dir) for p in get_weight_files(other_dir)]
|
|
if set(other_weights) != set(weight_rel_paths):
|
|
raise ValueError(f"Source {sources[i]} has different weight files")
|
|
|
|
# Create souped_dir
|
|
souped_dir = os.path.join(temp_dir, "souped")
|
|
# Copy first source (including its weights, which will be overwritten)
|
|
copy_local_to_local(first_source, souped_dir)
|
|
|
|
# Average weights
|
|
for rel_path in tqdm(weight_rel_paths, desc="Averaging weight files"):
|
|
all_paths = [os.path.join(st, rel_path) for st in source_temps]
|
|
file_path = all_paths[0]
|
|
souped_path = os.path.join(souped_dir, rel_path)
|
|
os.makedirs(os.path.dirname(souped_path), exist_ok=True)
|
|
|
|
if file_path.endswith(".safetensors"):
|
|
sum_state = load_file(file_path, device="cpu")
|
|
# Store original dtypes for each tensor
|
|
original_dtypes = {k: v.dtype for k, v in sum_state.items()}
|
|
# Upconvert to at least fp32 for accurate averaging
|
|
for k in sum_state:
|
|
if sum_state[k].dtype in (torch.float16, torch.bfloat16):
|
|
sum_state[k] = sum_state[k].to(torch.float32)
|
|
|
|
for other_path in all_paths[1:]:
|
|
other_state = load_file(other_path, device="cpu")
|
|
if set(sum_state.keys()) != set(other_state.keys()):
|
|
raise ValueError(f"Key mismatch in {rel_path}")
|
|
for k in sum_state:
|
|
# Upconvert other state to match sum_state dtype
|
|
if other_state[k].dtype in (torch.float16, torch.bfloat16):
|
|
other_state[k] = other_state[k].to(torch.float32)
|
|
sum_state[k] += other_state[k]
|
|
del other_state
|
|
|
|
n = len(all_paths)
|
|
for k in sum_state:
|
|
sum_state[k] /= n
|
|
# Cast back to original dtype
|
|
sum_state[k] = sum_state[k].to(original_dtypes[k])
|
|
save_file(sum_state, souped_path)
|
|
elif file_path.endswith(".bin"):
|
|
sum_state = torch.load(file_path, map_location="cpu")
|
|
# Store original dtypes for each tensor
|
|
original_dtypes = {k: v.dtype for k, v in sum_state.items()}
|
|
# Upconvert to at least fp32 for accurate averaging
|
|
for k in sum_state:
|
|
if sum_state[k].dtype in (torch.float16, torch.bfloat16):
|
|
sum_state[k] = sum_state[k].to(torch.float32)
|
|
|
|
for other_path in all_paths[1:]:
|
|
other_state = torch.load(other_path, map_location="cpu")
|
|
if set(sum_state.keys()) != set(other_state.keys()):
|
|
raise ValueError(f"Key mismatch in {rel_path}")
|
|
for k in sum_state:
|
|
# Upconvert other state to match sum_state dtype
|
|
if other_state[k].dtype in (torch.float16, torch.bfloat16):
|
|
other_state[k] = other_state[k].to(torch.float32)
|
|
sum_state[k] += other_state[k]
|
|
del other_state
|
|
|
|
n = len(all_paths)
|
|
for k in sum_state:
|
|
sum_state[k] /= n
|
|
# Cast back to original dtype
|
|
sum_state[k] = sum_state[k].to(original_dtypes[k])
|
|
torch.save(sum_state, souped_path)
|
|
else:
|
|
print(f"Skipping unknown weight file: {rel_path}")
|
|
continue
|
|
|
|
# Now copy souped_dir to dest_path
|
|
print("\nCopying souped model files to destination...")
|
|
if is_s3_path(dest_path):
|
|
dest_bucket, dest_prefix = parse_s3_path(dest_path)
|
|
copy_local_to_s3(souped_dir, dest_bucket, dest_prefix)
|
|
else:
|
|
copy_local_to_local(souped_dir, dest_path)
|
|
|
|
if final_architecture is None:
|
|
raise ValueError("Unable to determine the architecture of the prepared checkpoint.")
|
|
|
|
hf_model_id = HF_MODEL_IDS[final_architecture]
|
|
hf_base_url = f"https://huggingface.co/{hf_model_id}/resolve/main"
|
|
print(f"Using HuggingFace model: {hf_model_id}")
|
|
|
|
# Download tokenizer files from Hugging Face
|
|
print("\nDownloading tokenizer files from Hugging Face...")
|
|
|
|
if is_s3_path(dest_path):
|
|
# Download to temp directory first, then upload to S3
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
# Download files
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=6) as executor:
|
|
futures = [executor.submit(download_file_from_hf, filename, temp_dir, hf_base_url) for filename in TOKENIZER_FILES]
|
|
for future in concurrent.futures.as_completed(futures):
|
|
future.result()
|
|
|
|
# Upload to S3
|
|
dest_bucket, dest_prefix = parse_s3_path(dest_path)
|
|
upload_tasks = []
|
|
for filename in TOKENIZER_FILES:
|
|
local_path = os.path.join(temp_dir, filename)
|
|
s3_key = os.path.join(dest_prefix, filename)
|
|
upload_tasks.append((local_path, dest_bucket, s3_key))
|
|
|
|
print("Uploading tokenizer files to S3...")
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=6) as executor:
|
|
futures = [executor.submit(upload_file_to_s3, local_path, bucket, key) for local_path, bucket, key in upload_tasks]
|
|
for future in concurrent.futures.as_completed(futures):
|
|
future.result()
|
|
else:
|
|
# Download directly to destination
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=6) as executor:
|
|
futures = [executor.submit(download_file_from_hf, filename, dest_path, hf_base_url) for filename in TOKENIZER_FILES]
|
|
for future in concurrent.futures.as_completed(futures):
|
|
future.result()
|
|
|
|
print(f"\n✓ Successfully prepared checkpoint at {dest_path}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Prepare OlmOCR checkpoint for deployment",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog=__doc__.split("Usage:")[1], # Use the docstring for epilog
|
|
)
|
|
parser.add_argument("paths", nargs="+", help="One or more source paths followed by destination path (local or S3)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if len(args.paths) < 2:
|
|
parser.error("At least one source and one destination required")
|
|
|
|
sources = args.paths[:-1]
|
|
destination = args.paths[-1]
|
|
|
|
try:
|
|
prepare_checkpoints(sources, destination)
|
|
except Exception as e:
|
|
print(f"\n❌ Error: {e}")
|
|
return 1
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main())
|