mirror of
https://github.com/migatu/conjurer.git
synced 2026-07-14 13:34:40 +00:00
Fixes
This commit is contained in:
@@ -0,0 +1,193 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
from pathlib import Path
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
|
||||
|
||||
# =========================
|
||||
# Configuration constants
|
||||
# =========================
|
||||
|
||||
MAX_CHUNK_SIZE_BYTES = 50 * 1024 * 1024 # 100 MB
|
||||
OUTPUT_SUFFIX = "_chunk.txt"
|
||||
|
||||
|
||||
# =========================
|
||||
# Logging
|
||||
# =========================
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_default_output_dir(input_file: Path) -> Path:
|
||||
"""
|
||||
Return default output directory placed next to the input file.
|
||||
|
||||
Example:
|
||||
/data/big.txt -> /data/big/
|
||||
"""
|
||||
return input_file.with_suffix("")
|
||||
|
||||
|
||||
def write_chunk(output_dir: Path, chunk_index: int, lines: list[bytes], chunk_size: int) -> Path:
|
||||
"""
|
||||
Write one chunk file and return its path.
|
||||
"""
|
||||
output_file = output_dir / f"{chunk_index}{OUTPUT_SUFFIX}"
|
||||
|
||||
with output_file.open("wb") as file:
|
||||
file.writelines(lines)
|
||||
|
||||
actual_size = output_file.stat().st_size
|
||||
|
||||
if actual_size != chunk_size:
|
||||
logger.warning(
|
||||
"Chunk size mismatch for %s: expected %d bytes, got %d bytes",
|
||||
output_file,
|
||||
chunk_size,
|
||||
actual_size,
|
||||
)
|
||||
|
||||
return output_file
|
||||
|
||||
|
||||
def split_txt_file(input_file: Path, output_dir: Path | None = None) -> None:
|
||||
"""
|
||||
Split a text-like file into smaller files limited by MAX_CHUNK_SIZE_BYTES.
|
||||
|
||||
The file is processed in binary mode, so invalid UTF-8 or mixed encodings
|
||||
do not break the split. Lines are never split.
|
||||
"""
|
||||
if not input_file.exists():
|
||||
raise FileNotFoundError(f"Input file does not exist: {input_file}")
|
||||
|
||||
if not input_file.is_file():
|
||||
raise ValueError(f"Input path is not a file: {input_file}")
|
||||
|
||||
if MAX_CHUNK_SIZE_BYTES <= 0:
|
||||
raise ValueError("MAX_CHUNK_SIZE_BYTES must be greater than zero")
|
||||
|
||||
if output_dir is None:
|
||||
output_dir = get_default_output_dir(input_file)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger.info("Input file: %s", input_file)
|
||||
logger.info("Output directory: %s", output_dir)
|
||||
logger.info("Max chunk size: %d bytes", MAX_CHUNK_SIZE_BYTES)
|
||||
|
||||
chunk_index = 0
|
||||
current_lines: list[bytes] = []
|
||||
current_size = 0
|
||||
total_lines = 0
|
||||
total_bytes = 0
|
||||
|
||||
with input_file.open("rb") as input_handle:
|
||||
for line_number, line in enumerate(input_handle, start=1):
|
||||
line_size = len(line)
|
||||
total_lines = line_number
|
||||
total_bytes += line_size
|
||||
|
||||
if line_size > MAX_CHUNK_SIZE_BYTES:
|
||||
logger.warning(
|
||||
"Line %d is larger than max chunk size: %d bytes > %d bytes. "
|
||||
"It will be written as a separate chunk.",
|
||||
line_number,
|
||||
line_size,
|
||||
MAX_CHUNK_SIZE_BYTES,
|
||||
)
|
||||
|
||||
if current_lines and current_size + line_size > MAX_CHUNK_SIZE_BYTES:
|
||||
output_file = write_chunk(
|
||||
output_dir=output_dir,
|
||||
chunk_index=chunk_index,
|
||||
lines=current_lines,
|
||||
chunk_size=current_size,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Written chunk %d: %s (%d bytes)",
|
||||
chunk_index,
|
||||
output_file,
|
||||
current_size,
|
||||
)
|
||||
|
||||
chunk_index += 1
|
||||
current_lines = []
|
||||
current_size = 0
|
||||
|
||||
current_lines.append(line)
|
||||
current_size += line_size
|
||||
|
||||
if current_lines:
|
||||
output_file = write_chunk(
|
||||
output_dir=output_dir,
|
||||
chunk_index=chunk_index,
|
||||
lines=current_lines,
|
||||
chunk_size=current_size,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Written chunk %d: %s (%d bytes)",
|
||||
chunk_index,
|
||||
output_file,
|
||||
current_size,
|
||||
)
|
||||
|
||||
chunk_count = chunk_index + 1
|
||||
else:
|
||||
logger.info("Input file is empty. No chunks were created.")
|
||||
chunk_count = 0
|
||||
|
||||
logger.info("Total lines processed: %d", total_lines)
|
||||
logger.info("Total bytes processed: %d", total_bytes)
|
||||
logger.info("Chunks created: %d", chunk_count)
|
||||
logger.info("Done.")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Split a large TXT-like file into smaller line-safe chunks."
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"input_file",
|
||||
type=Path,
|
||||
help="Path to input TXT file.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--output-dir",
|
||||
type=Path,
|
||||
default=None,
|
||||
help=(
|
||||
"Optional output directory. "
|
||||
"By default, a directory named after the input file is created next to it."
|
||||
),
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
|
||||
try:
|
||||
split_txt_file(args.input_file, args.output_dir)
|
||||
return 0
|
||||
except Exception:
|
||||
logger.exception("Failed to split file")
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
||||
@@ -19,7 +19,7 @@ from requests import ConnectTimeout, Timeout
|
||||
import lib_paths
|
||||
|
||||
SCR_DATABASE_PATH = os.getenv("CONJURER_LIBRARIAN_DB_PATH", r"C:\\Database\\chunks\\")
|
||||
SCR_FILENAME = os.getenv("CONJURER_LIBRARIAN_SCRAPE_CHUNK", "40_chunk.txt")
|
||||
SCR_FILENAME = os.getenv("CONJURER_LIBRARIAN_SCRAPE_CHUNK", "44_chunk.txt")
|
||||
SCR_ENCODING = os.getenv("CONJURER_ENCODING", "utf-8")
|
||||
|
||||
WORK_Q = Queue()
|
||||
|
||||
@@ -16,7 +16,7 @@ services:
|
||||
volumes:
|
||||
# Local DOI chunk database (0_chunk.txt ... N_chunk.txt). Read-only:
|
||||
# the search workers only read it (open mode "r").
|
||||
- /srv/librarian/doi:/doi:ro
|
||||
- /mnt/conjurer_swap/librarian_data/base_it1:/doi:ro
|
||||
# Runtime JSON state (cr_results/rr_results/not_in_db/s_results) -
|
||||
# seeded on first run, persisted here across restarts.
|
||||
- /srv/librarian/state:/lib_temp_files
|
||||
|
||||
Reference in New Issue
Block a user