Function main

src/docquery/cmd/scan.py:52–103

(args)

def main(args):
    paths = []
    if pathlib.Path(args.path).is_dir():
        for root, dirs, files in os.walk(args.path):
            for fname in files:
                if (pathlib.Path(root) / fname).is_dir():
                    continue
                paths.append(pathlib.Path(root) / fname)
    else:
        paths.append(args.path)

    docs = []
    for p in paths:
        try:
            log.info(f"Loading {p}")
            docs.append((p, load_document(str(p), ocr_reader=args.ocr, use_embedded_text=args.use_embedded_text)))
        except UnsupportedDocument as e:
            log.warning(f"Cannot load {p}: {e}. Skipping...")

    log.info(f"Done loading {len(docs)} file(s).")
    if not docs:
        return

    log.info("Loading pipelines.")

    nlp = pipeline("document-question-answering", model=args.checkpoint)
    if args.classify:
        classify = pipeline("document-classification", model=args.classify_checkpoint)

    log.info("Ready to start evaluating!")

    max_fname_len = max(len(str(p)) for (p, _) in docs)
    max_question_len = max(len(q) for q in args.questions) if len(args.questions) > 0 else 0
    for i, (p, d) in enumerate(docs):
        if i > 0 and len(args.questions) > 1:
            print("")

        if args.classify:
            cls = classify(**d.context)[0]
            print(f"{str(p):<{max_fname_len}} Document Type: {cls['label']}")

        for q in args.questions:
            try:
                response = nlp(question=q, **d.context)
                if isinstance(response, list):
                    response = response[0] if len(response) > 0 else None
            except Exception:
                log.error(f"Failed while processing {str(p)} on question: '{q}'")
                raise

            answer = response["answer"] if response is not None else "NULL"
            print(f"{str(p):<{max_fname_len}} {q:<{max_question_len}}: {answer}")