Method _handle_documents

app/src/token_count.py:207–259

(
            self,
            client: botocore.client.BaseClient,
            input_id: str,
            decoder
    )

        return result

    def _handle_documents(
            self,
            client: botocore.client.BaseClient,
            input_id: str,
            decoder
    ) -> InputResult:
        # download doc
        input_uri = urlparse(
            os.path.join(
                self._input_base_uri, f"{input_id}.json.gz"
            )
        )
        dl_status, input_buffer = self._dload_file(input_uri, client=client)

        # check if download was successful
        if not dl_status.is_success:
            return InputResult(
                is_success=False, msg=dl_status.msg, input_id=input_id
            )

        num_docs = 0
        total_tokens = 0
        token_counts = []

        try:
            with gzip.open(input_buffer, mode="rb") as in_fh:
                for idx, obj in enumerate(in_fh):
                    record = decoder.decode(obj)

                    # tokenize
                    num_tokens = len(
                        TOKENIZER.encode(record.raw_content).tokens
                    )
                    token_counts.append((idx, num_tokens))

                    total_tokens += num_tokens
                    num_docs += 1

        except Exception as e:
            msg = (
                f"__DECODE_ENCODE_FAIL__ {input_id}: "
                f"caught exception {e.__class__.__name__}: {e}"
            )
            return InputResult(is_success=False, msg=msg, input_id=input_id)

        return InputResult(
            is_success=True,
            msg="",
            input_id=input_id,
            num_docs=num_docs,
            num_tokens=total_tokens,
            token_counts=token_counts
        )

    def run(self):
        # init logging