From 6ca5e4d48f3d02ed2323514281bddcf4b482678c Mon Sep 17 00:00:00 2001 From: root Date: Thu, 9 Jul 2026 14:55:35 +0200 Subject: [PATCH] Fixes --- conjurer_librarian/cut_pieces.py | 193 +++++++++++++++++++++++++++++++ conjurer_librarian/scrape_bot.py | 2 +- docker/compose.librarian.yaml | 2 +- 3 files changed, 195 insertions(+), 2 deletions(-) create mode 100644 conjurer_librarian/cut_pieces.py diff --git a/conjurer_librarian/cut_pieces.py b/conjurer_librarian/cut_pieces.py new file mode 100644 index 0000000..28e13c1 --- /dev/null +++ b/conjurer_librarian/cut_pieces.py @@ -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()) + diff --git a/conjurer_librarian/scrape_bot.py b/conjurer_librarian/scrape_bot.py index 73e8e9e..cd7fdef 100644 --- a/conjurer_librarian/scrape_bot.py +++ b/conjurer_librarian/scrape_bot.py @@ -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() diff --git a/docker/compose.librarian.yaml b/docker/compose.librarian.yaml index 07b2c6f..82c99b5 100644 --- a/docker/compose.librarian.yaml +++ b/docker/compose.librarian.yaml @@ -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