chore: import upstream snapshot with attribution
Main / Python 3.11 - Docs (push) Waiting to run
Main / Python 3.11 - Build (push) Waiting to run
Main / Python 3.11 - Lint (push) Waiting to run
Main / Python 3.11 - Style (push) Waiting to run
Main / Python 3.11 - Test (push) Waiting to run
Main / GPU CI (push) Blocked by required conditions
Main / Release (push) Blocked by required conditions
Main / Build and Push Docker Images (push) Blocked by required conditions
Main / Python 3.11 - Docs (push) Waiting to run
Main / Python 3.11 - Build (push) Waiting to run
Main / Python 3.11 - Lint (push) Waiting to run
Main / Python 3.11 - Style (push) Waiting to run
Main / Python 3.11 - Test (push) Waiting to run
Main / GPU CI (push) Blocked by required conditions
Main / Release (push) Blocked by required conditions
Main / Build and Push Docker Images (push) Blocked by required conditions
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,18 @@
|
||||
.git
|
||||
.github
|
||||
.mypy_cache
|
||||
.pytest_cache
|
||||
.venv
|
||||
__pycache__
|
||||
*.egg-info
|
||||
|
||||
|
||||
# Let's not copy any bash scripts from the scripts folder over, otherwise trashing the docker image too much with recent changes
|
||||
scripts/*.sh
|
||||
scripts/**/*.sh
|
||||
|
||||
# Nor copy any olmocr bench files
|
||||
olmOCR-bench/
|
||||
olmOCR-bench*/
|
||||
html_templates*/
|
||||
olmocr-synthmix-*/
|
||||
@@ -0,0 +1,168 @@
|
||||
# Contributing
|
||||
|
||||
Thanks for considering contributing! Please read this document to learn the various ways you can contribute to this project and how to go about doing it.
|
||||
|
||||
## Bug reports and feature requests
|
||||
|
||||
### Did you find a bug?
|
||||
|
||||
First, do [a quick search](https://github.com/allenai/olmocrissues) to see whether your issue has already been reported.
|
||||
If your issue has already been reported, please comment on the existing issue.
|
||||
|
||||
Otherwise, open [a new GitHub issue](https://github.com/allenai/olmocrissues). Be sure to include a clear title
|
||||
and description. The description should include as much relevant information as possible. The description should
|
||||
explain how to reproduce the erroneous behavior as well as the behavior you expect to see. Ideally you would include a
|
||||
code sample or an executable test case demonstrating the expected behavior.
|
||||
|
||||
### Do you have a suggestion for an enhancement or new feature?
|
||||
|
||||
We use GitHub issues to track feature requests. Before you create a feature request:
|
||||
|
||||
* Make sure you have a clear idea of the enhancement you would like. If you have a vague idea, consider discussing
|
||||
it first on a GitHub issue.
|
||||
* Check the documentation to make sure your feature does not already exist.
|
||||
* Do [a quick search](https://github.com/allenai/olmocrissues) to see whether your feature has already been suggested.
|
||||
|
||||
When creating your request, please:
|
||||
|
||||
* Provide a clear title and description.
|
||||
* Explain why the enhancement would be useful. It may be helpful to highlight the feature in other libraries.
|
||||
* Include code examples to demonstrate how the enhancement would be used.
|
||||
|
||||
## Making a pull request
|
||||
|
||||
When you're ready to contribute code to address an open issue, please follow these guidelines to help us be able to review your pull request (PR) quickly.
|
||||
|
||||
1. **Initial setup** (only do this once)
|
||||
|
||||
<details><summary>Expand details 👇</summary><br/>
|
||||
|
||||
If you haven't already done so, please [fork](https://help.github.com/en/enterprise/2.13/user/articles/fork-a-repo) this repository on GitHub.
|
||||
|
||||
Then clone your fork locally with
|
||||
|
||||
git clone https://github.com/USERNAME/olmocrgit
|
||||
|
||||
or
|
||||
|
||||
git clone git@github.com:USERNAME/olmocrgit
|
||||
|
||||
At this point the local clone of your fork only knows that it came from *your* repo, github.com/USERNAME/olmocrgit, but doesn't know anything the *main* repo, [https://github.com/allenai/oolmocrit](https://github.com/allenai/ololmocrYou can see this by running
|
||||
|
||||
git remote -v
|
||||
|
||||
which will output something like this:
|
||||
|
||||
origin https://github.com/USERNAME/olmocrgit (fetch)
|
||||
origin https://github.com/USERNAME/olmocrgit (push)
|
||||
|
||||
This means that your local clone can only track changes from your fork, but not from the main repo, and so you won't be able to keep your fork up-to-date with the main repo over time. Therefore you'll need to add another "remote" to your clone that points to [https://github.com/allenai/olmocrgit](https://github.com/allenai/oolmocr To do this, run the following:
|
||||
|
||||
git remote add upstream https://github.com/allenai/olmocrgit
|
||||
|
||||
Now if you do `git remote -v` again, you'll see
|
||||
|
||||
origin https://github.com/USERNAME/olmocrgit (fetch)
|
||||
origin https://github.com/USERNAME/olmocrgit (push)
|
||||
upstream https://github.com/allenai/olmocrgit (fetch)
|
||||
upstream https://github.com/allenai/olmocrgit (push)
|
||||
|
||||
Finally, you'll need to create a Python 3 virtual environment suitable for working on this project. There a number of tools out there that making working with virtual environments easier.
|
||||
The most direct way is with the [`venv` module](https://docs.python.org/3.7/library/venv.html) in the standard library, but if you're new to Python or you don't already have a recent Python 3 version installed on your machine,
|
||||
we recommend [Miniconda](https://docs.conda.io/en/latest/miniconda.html).
|
||||
|
||||
On Mac, for example, you can install Miniconda with [Homebrew](https://brew.sh/):
|
||||
|
||||
brew install miniconda
|
||||
|
||||
Then you can create and activate a new Python environment by running:
|
||||
|
||||
conda create -n olmocrpython=3.9
|
||||
conda activate olmocr
|
||||
|
||||
Once your virtual environment is activated, you can install your local clone in "editable mode" with
|
||||
|
||||
pip install -U pip setuptools wheel
|
||||
pip install -e .[dev]
|
||||
|
||||
The "editable mode" comes from the `-e` argument to `pip`, and essential just creates a symbolic link from the site-packages directory of your virtual environment to the source code in your local clone. That way any changes you make will be immediately reflected in your virtual environment.
|
||||
|
||||
</details>
|
||||
|
||||
2. **Ensure your fork is up-to-date**
|
||||
|
||||
<details><summary>Expand details 👇</summary><br/>
|
||||
|
||||
Once you've added an "upstream" remote pointing to [https://github.com/allenai/python-package-temlate.git](https://github.com/allenai/olmocr, keeping your fork up-to-date is easy:
|
||||
|
||||
git checkout main # if not already on main
|
||||
git pull --rebase upstream main
|
||||
git push
|
||||
|
||||
</details>
|
||||
|
||||
3. **Create a new branch to work on your fix or enhancement**
|
||||
|
||||
<details><summary>Expand details 👇</summary><br/>
|
||||
|
||||
Committing directly to the main branch of your fork is not recommended. It will be easier to keep your fork clean if you work on a separate branch for each contribution you intend to make.
|
||||
|
||||
You can create a new branch with
|
||||
|
||||
# replace BRANCH with whatever name you want to give it
|
||||
git checkout -b BRANCH
|
||||
git push -u origin BRANCH
|
||||
|
||||
</details>
|
||||
|
||||
4. **Test your changes**
|
||||
|
||||
<details><summary>Expand details 👇</summary><br/>
|
||||
|
||||
Our continuous integration (CI) testing runs [a number of checks](https://github.com/allenai/olmocractions) for each pull request on [GitHub Actions](https://github.com/features/actions). You can run most of these tests locally, which is something you should do *before* opening a PR to help speed up the review process and make it easier for us.
|
||||
|
||||
First, you should run [`isort`](https://github.com/PyCQA/isort) and [`black`](https://github.com/psf/black) to make sure you code is formatted consistently.
|
||||
Many IDEs support code formatters as plugins, so you may be able to setup isort and black to run automatically everytime you save.
|
||||
For example, [`black.vim`](https://github.com/psf/black/tree/master/plugin) will give you this functionality in Vim. But both `isort` and `black` are also easy to run directly from the command line.
|
||||
Just run this from the root of your clone:
|
||||
|
||||
isort .
|
||||
black .
|
||||
|
||||
Our CI also uses [`ruff`](https://github.com/astral-sh/ruff) to lint the code base and [`mypy`](http://mypy-lang.org/) for type-checking. You should run both of these next with
|
||||
|
||||
ruff check .
|
||||
|
||||
and
|
||||
|
||||
mypy .
|
||||
|
||||
We also strive to maintain high test coverage, so most contributions should include additions to [the unit tests](https://github.com/allenai/olmocrtree/main/tests). These tests are run with [`pytest`](https://docs.pytest.org/en/latest/), which you can use to locally run any test modules that you've added or changed.
|
||||
|
||||
For example, if you've fixed a bug in `olmocra/b.py`, you can run the tests specific to that module with
|
||||
|
||||
pytest -v tests/a/b_test.py
|
||||
|
||||
If your contribution involves additions to any public part of the API, we require that you write docstrings
|
||||
for each function, method, class, or module that you add.
|
||||
See the [Writing docstrings](#writing-docstrings) section below for details on the syntax.
|
||||
You should test to make sure the API documentation can build without errors by running
|
||||
|
||||
make docs
|
||||
|
||||
If the build fails, it's most likely due to small formatting issues. If the error message isn't clear, feel free to comment on this in your pull request.
|
||||
|
||||
And finally, please update the [CHANGELOG](https://github.com/allenai/olmocrblob/main/CHANGELOG.md) with notes on your contribution in the "Unreleased" section at the top.
|
||||
|
||||
After all of the above checks have passed, you can now open [a new GitHub pull request](https://github.com/allenai/olmocrpulls).
|
||||
Make sure you have a clear description of the problem and the solution, and include a link to relevant issues.
|
||||
|
||||
We look forward to reviewing your PR!
|
||||
|
||||
</details>
|
||||
|
||||
### Writing docstrings
|
||||
|
||||
We use [Sphinx](https://www.sphinx-doc.org/en/master/index.html) to build our API docs, which automatically parses all docstrings
|
||||
of public classes and methods using the [autodoc](https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html) extension.
|
||||
Please refer to autoc's documentation to learn about the docstring syntax.
|
||||
@@ -0,0 +1,46 @@
|
||||
name: 🐛 Bug Report
|
||||
description: Create a report to help us reproduce and fix the bug
|
||||
labels: 'bug'
|
||||
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: >
|
||||
#### Before submitting a bug, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/allenai/olmocr/issues?q=is%3Aissue+sort%3Acreated-desc+).
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: 🐛 Describe the bug
|
||||
description: |
|
||||
Please provide a clear and concise description of what the bug is.
|
||||
|
||||
If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
|
||||
|
||||
```python
|
||||
# All necessary imports at the beginning
|
||||
import olmocr
|
||||
|
||||
# A succinct reproducing example trimmed down to the essential parts:
|
||||
assert False is True, "Oh no!"
|
||||
```
|
||||
|
||||
If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.
|
||||
|
||||
Please also paste or describe the results you observe instead of the expected results. If you observe an error, please paste the error message including the **full** traceback of the exception. It may be relevant to wrap error messages in ```` ```triple quotes blocks``` ````.
|
||||
placeholder: |
|
||||
A clear and concise description of what the bug is.
|
||||
validations:
|
||||
required: true
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Versions
|
||||
description: |
|
||||
Please run the following and paste the output below.
|
||||
```sh
|
||||
python --version && pip freeze
|
||||
```
|
||||
validations:
|
||||
required: true
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: >
|
||||
Thanks for contributing 🎉!
|
||||
@@ -0,0 +1,21 @@
|
||||
name: 📚 Documentation
|
||||
description: Report an issue related to https://olmocr.readthedocs.io/latest
|
||||
labels: 'documentation'
|
||||
|
||||
body:
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: 📚 The doc issue
|
||||
description: >
|
||||
A clear and concise description of what content in https://olmocr.readthedocs.io/latest is an issue.
|
||||
validations:
|
||||
required: true
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Suggest a potential alternative/fix
|
||||
description: >
|
||||
Tell us how we could improve the documentation in this regard.
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: >
|
||||
Thanks for contributing 🎉!
|
||||
@@ -0,0 +1,26 @@
|
||||
name: 🚀 Feature request
|
||||
description: Submit a proposal/request for a new feature
|
||||
labels: 'feature request'
|
||||
|
||||
body:
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: 🚀 The feature, motivation and pitch
|
||||
description: >
|
||||
A clear and concise description of the feature proposal. Please outline the motivation for the proposal. Is your feature request related to a specific problem? e.g., *"I'm working on X and would like Y to be possible"*. If this is related to another GitHub issue, please link here too.
|
||||
validations:
|
||||
required: true
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Alternatives
|
||||
description: >
|
||||
A description of any alternative solutions or features you've considered, if any.
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Additional context
|
||||
description: >
|
||||
Add any other context or screenshots about the feature request.
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: >
|
||||
Thanks for contributing 🎉!
|
||||
@@ -0,0 +1,58 @@
|
||||
name: Python virtualenv
|
||||
description: Set up a Python virtual environment with caching
|
||||
inputs:
|
||||
python-version:
|
||||
description: The Python version to use
|
||||
required: true
|
||||
cache-prefix:
|
||||
description: Update this to invalidate the cache
|
||||
required: true
|
||||
default: v0
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
- name: Setup Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: ${{ inputs.python-version }}
|
||||
|
||||
- shell: bash
|
||||
run: |
|
||||
# Install prerequisites.
|
||||
pip install --upgrade pip setuptools wheel virtualenv
|
||||
|
||||
- shell: bash
|
||||
run: |
|
||||
# Get the exact Python version to use in the cache key.
|
||||
echo "PYTHON_VERSION=$(python --version)" >> $GITHUB_ENV
|
||||
|
||||
- uses: actions/cache@v3
|
||||
id: virtualenv-cache
|
||||
with:
|
||||
path: .venv
|
||||
key: ${{ inputs.cache-prefix }}-${{ runner.os }}-${{ env.PYTHON_VERSION }}-${{ hashFiles('pyproject.toml') }}
|
||||
|
||||
- if: steps.virtualenv-cache.outputs.cache-hit != 'true'
|
||||
shell: bash
|
||||
run: |
|
||||
# Set up virtual environment without cache hit.
|
||||
test -d .venv || virtualenv -p $(which python) --copies --reset-app-data .venv
|
||||
. .venv/bin/activate
|
||||
pip install -e .[dev]
|
||||
pip install -e .[bench]
|
||||
|
||||
- if: steps.virtualenv-cache.outputs.cache-hit == 'true'
|
||||
shell: bash
|
||||
run: |
|
||||
# Set up virtual environment from cache hit.
|
||||
. .venv/bin/activate
|
||||
pip install --no-deps -e .[dev]
|
||||
pip install --no-deps -e .[bench]
|
||||
|
||||
- shell: bash
|
||||
run: |
|
||||
# Show environment info.
|
||||
. .venv/bin/activate
|
||||
echo "✓ Installed $(python --version) virtual environment to $(which python)"
|
||||
echo "Packages:"
|
||||
pip freeze
|
||||
@@ -0,0 +1,11 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "pip"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "daily"
|
||||
open-pull-requests-limit: 10
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "daily"
|
||||
@@ -0,0 +1,18 @@
|
||||
<!-- To ensure we can review your pull request promptly please complete this template entirely. -->
|
||||
|
||||
<!-- Please reference the issue number here. You can replace "Fixes" with "Closes" if it makes more sense. -->
|
||||
Fixes #
|
||||
|
||||
Changes proposed in this pull request:
|
||||
<!-- Please list all changes/additions here. -->
|
||||
-
|
||||
|
||||
## Before submitting
|
||||
|
||||
<!-- Please complete this checklist BEFORE submitting your PR to speed along the review process. -->
|
||||
- [ ] I've read and followed all steps in the [Making a pull request](https://github.com/allenai/olmocr/blob/main/.github/CONTRIBUTING.md#making-a-pull-request)
|
||||
section of the `CONTRIBUTING` docs.
|
||||
- [ ] I've updated or added any relevant docstrings following the syntax described in the
|
||||
[Writing docstrings](https://github.com/allenai/olmocr/blob/main/.github/CONTRIBUTING.md#writing-docstrings) section of the `CONTRIBUTING` docs.
|
||||
- [ ] If this PR fixes a bug, I've added a test that will fail without my fix.
|
||||
- [ ] If this PR adds a new feature, I've added tests that sufficiently cover my new functionality.
|
||||
@@ -0,0 +1,321 @@
|
||||
name: Main
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
tags:
|
||||
- "v*.*.*"
|
||||
|
||||
env:
|
||||
# Change this to invalidate existing cache.
|
||||
CACHE_PREFIX: v0
|
||||
PYTHONPATH: ./
|
||||
|
||||
jobs:
|
||||
checks:
|
||||
name: Python ${{ matrix.python }} - ${{ matrix.task.name }}
|
||||
runs-on: large-olmocr-runner
|
||||
timeout-minutes: 15
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
python: ["3.11"]
|
||||
task:
|
||||
- name: Test
|
||||
run: |
|
||||
playwright install chromium
|
||||
pytest -v --color=yes -m "not nonci" tests/
|
||||
|
||||
include:
|
||||
- python: "3.11"
|
||||
task:
|
||||
name: Lint
|
||||
run: ruff check .
|
||||
|
||||
# Removing mypy for now, as it isn't handling async things correctly and crashing
|
||||
# - python: "3.11"
|
||||
# task:
|
||||
# name: Type check
|
||||
# run: mypy .
|
||||
|
||||
- python: "3.11"
|
||||
task:
|
||||
name: Build
|
||||
run: |
|
||||
python -m build
|
||||
|
||||
- python: "3.11"
|
||||
task:
|
||||
name: Style
|
||||
run: |
|
||||
isort --check .
|
||||
black --check .
|
||||
|
||||
- python: "3.11"
|
||||
task:
|
||||
name: Docs
|
||||
run: cd docs && make html
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
|
||||
- name: Install system dependencies
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install -y --no-install-recommends poppler-utils
|
||||
|
||||
- name: Setup Python environment
|
||||
uses: ./.github/actions/setup-venv
|
||||
with:
|
||||
python-version: ${{ matrix.python }}
|
||||
cache-prefix: ${{ env.CACHE_PREFIX }}
|
||||
|
||||
- name: Restore mypy cache
|
||||
if: matrix.task.name == 'Type check'
|
||||
uses: actions/cache@v3
|
||||
with:
|
||||
path: .mypy_cache
|
||||
key: mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt') }}-${{ github.ref }}-${{ github.sha }}
|
||||
restore-keys: |
|
||||
mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt') }}-${{ github.ref }}
|
||||
mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt') }}
|
||||
|
||||
- name: ${{ matrix.task.name }}
|
||||
run: |
|
||||
. .venv/bin/activate
|
||||
${{ matrix.task.run }}
|
||||
|
||||
- name: Upload package distribution files
|
||||
if: matrix.task.name == 'Build'
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: package
|
||||
path: dist
|
||||
|
||||
- name: Clean up
|
||||
if: always()
|
||||
run: |
|
||||
. .venv/bin/activate
|
||||
pip uninstall -y olmocr
|
||||
|
||||
gpu_checks:
|
||||
name: GPU CI
|
||||
runs-on: large-olmocr-runner
|
||||
timeout-minutes: 15
|
||||
needs: [checks]
|
||||
env:
|
||||
BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
|
||||
BEAKER_IMAGE: jakep/olmocr-gpu-ci
|
||||
BEAKER_BUDGET: ai2/oe-base
|
||||
BEAKER_WORKSPACE: ai2/olmocr
|
||||
steps:
|
||||
- name: Determine current commit SHA (pull request)
|
||||
if: github.event_name == 'pull_request'
|
||||
run: |
|
||||
echo "COMMIT_SHA=${{ github.event.pull_request.head.sha }}" >> $GITHUB_ENV
|
||||
|
||||
- name: Determine current commit SHA (push)
|
||||
if: github.event_name != 'pull_request'
|
||||
run: |
|
||||
echo "COMMIT_SHA=$GITHUB_SHA" >> $GITHUB_ENV
|
||||
|
||||
- name: GPU Tests
|
||||
uses: allenai/beaker-run-action@v1.2
|
||||
if: env.BEAKER_TOKEN != ''
|
||||
with:
|
||||
spec: |
|
||||
version: v2
|
||||
description: GPU Tests
|
||||
budget: ${{ env.BEAKER_BUDGET }}
|
||||
tasks:
|
||||
- name: tests
|
||||
image:
|
||||
beaker: ${{ env.BEAKER_IMAGE }}
|
||||
context:
|
||||
priority: normal
|
||||
preemptible: true
|
||||
resources:
|
||||
gpuCount: 1
|
||||
constraints:
|
||||
cluster:
|
||||
- ai2/jupiter-cirrascale-2
|
||||
- ai2/neptune-cirrascale
|
||||
- ai2/saturn-cirrascale
|
||||
- ai2/ceres-cirrascale
|
||||
envVars:
|
||||
- name: GIT_REVISION
|
||||
value: ${{ env.COMMIT_SHA }}
|
||||
entrypoint: ["/bin/bash"]
|
||||
command: ["./gpu-ci-script.sh"]
|
||||
result:
|
||||
path: /unused
|
||||
token: ${{ env.BEAKER_TOKEN }}
|
||||
workspace: ${{ env.BEAKER_WORKSPACE }}
|
||||
|
||||
release:
|
||||
name: Release
|
||||
runs-on: large-olmocr-runner
|
||||
needs: [checks, gpu_checks]
|
||||
if: startsWith(github.ref, 'refs/tags/')
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Setup Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: "3.11"
|
||||
|
||||
- name: Install requirements
|
||||
run: |
|
||||
pip install --upgrade pip setuptools wheel build
|
||||
pip install -e .[dev]
|
||||
pip install -e .[bench]
|
||||
|
||||
- name: Prepare environment
|
||||
run: |
|
||||
echo "RELEASE_VERSION=${GITHUB_REF#refs/tags/v}" >> $GITHUB_ENV
|
||||
echo "TAG=${GITHUB_REF#refs/tags/}" >> $GITHUB_ENV
|
||||
|
||||
- name: Download package distribution files
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: package
|
||||
path: dist
|
||||
|
||||
- name: Generate release notes
|
||||
run: |
|
||||
python scripts/release_notes.py > ${{ github.workspace }}-RELEASE_NOTES.md
|
||||
|
||||
- name: Publish package to PyPI
|
||||
run: |
|
||||
twine upload -u '${{ secrets.PYPI_USERNAME }}' -p '${{ secrets.PYPI_PASSWORD }}' dist/*
|
||||
|
||||
- name: Publish GitHub release
|
||||
uses: softprops/action-gh-release@v1
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
with:
|
||||
body_path: ${{ github.workspace }}-RELEASE_NOTES.md
|
||||
prerelease: ${{ contains(env.TAG, 'rc') }}
|
||||
files: |
|
||||
dist/*
|
||||
|
||||
docker-build:
|
||||
name: Build and Push Docker Images
|
||||
runs-on: large-olmocr-runner
|
||||
needs: [release]
|
||||
if: startsWith(github.ref, 'refs/tags/')
|
||||
permissions:
|
||||
contents: read
|
||||
packages: write
|
||||
env:
|
||||
REGISTRY: docker.io
|
||||
IMAGE_NAME: alleninstituteforai/olmocr
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Free up disk space
|
||||
run: |
|
||||
sudo rm -rf /usr/share/dotnet \
|
||||
/usr/local/lib/android \
|
||||
/opt/ghc \
|
||||
/opt/hostedtoolcache/CodeQL
|
||||
sudo docker system prune -af
|
||||
sudo apt-get -y autoremove
|
||||
sudo apt-get -y autoclean
|
||||
df -h
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
|
||||
- name: Log in to Docker Hub
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
registry: ${{ env.REGISTRY }}
|
||||
username: ${{ secrets.DOCKER_USERNAME }}
|
||||
password: ${{ secrets.DOCKER_PASSWORD }}
|
||||
|
||||
- name: Extract metadata for base image
|
||||
id: meta-base
|
||||
uses: docker/metadata-action@v5
|
||||
with:
|
||||
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
|
||||
tags: |
|
||||
type=ref,event=tag
|
||||
type=semver,pattern={{version}}
|
||||
type=semver,pattern={{major}}.{{minor}}
|
||||
flavor: |
|
||||
latest=true
|
||||
|
||||
- name: Build and push base Docker image
|
||||
uses: docker/build-push-action@v5
|
||||
with:
|
||||
context: .
|
||||
file: ./Dockerfile
|
||||
push: true
|
||||
tags: ${{ steps.meta-base.outputs.tags }}
|
||||
labels: ${{ steps.meta-base.outputs.labels }}
|
||||
platforms: linux/amd64
|
||||
outputs: type=registry
|
||||
no-cache: true
|
||||
|
||||
- name: Extract metadata for image with model
|
||||
id: meta-with-model
|
||||
uses: docker/metadata-action@v5
|
||||
with:
|
||||
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
|
||||
tags: |
|
||||
type=ref,event=tag,suffix=-with-model
|
||||
type=semver,pattern={{version}}-with-model
|
||||
type=semver,pattern={{major}}.{{minor}}-with-model
|
||||
flavor: |
|
||||
latest=auto
|
||||
suffix=-with-model,onlatest=true
|
||||
|
||||
- name: Build and push Docker image with model
|
||||
uses: docker/build-push-action@v5
|
||||
with:
|
||||
context: .
|
||||
file: ./Dockerfile.with-model
|
||||
push: true
|
||||
tags: ${{ steps.meta-with-model.outputs.tags }}
|
||||
labels: ${{ steps.meta-with-model.outputs.labels }}
|
||||
platforms: linux/amd64
|
||||
outputs: type=registry
|
||||
|
||||
# jakep: push to beaker can't work because of limitted disk space on these runners
|
||||
# jakep: (you can try by setting load: true above, but you'll need a larger runner)
|
||||
# - name: Setup Beaker CLI
|
||||
# uses: allenai/setup-beaker@v2
|
||||
# with:
|
||||
# token: ${{ secrets.BEAKER_TOKEN }}
|
||||
# version: latest
|
||||
# - name: Debug Docker images
|
||||
# run: docker images
|
||||
|
||||
# - name: Push to Beaker
|
||||
# env:
|
||||
# BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
|
||||
# run: |
|
||||
# VERSION=${{ steps.meta.outputs.version }}
|
||||
# beaker image create \
|
||||
# --name "olmocr-inference-$VERSION" \
|
||||
# --workspace ai2/olmocr \
|
||||
# alleninstituteforai/olmocr:$VERSION
|
||||
|
||||
- name: Clean up after build
|
||||
if: always()
|
||||
run: |
|
||||
docker system prune -af --volumes
|
||||
df -h
|
||||
@@ -0,0 +1,29 @@
|
||||
name: PR Checks
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- 'olmocr/**'
|
||||
|
||||
jobs:
|
||||
changelog:
|
||||
name: CHANGELOG
|
||||
runs-on: ubuntu-latest
|
||||
if: github.event_name == 'pull_request'
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Check that CHANGELOG has been updated
|
||||
run: |
|
||||
# If this step fails, this means you haven't updated the CHANGELOG.md
|
||||
# file with notes on your contribution.
|
||||
git diff --name-only $(git merge-base origin/main HEAD) | grep '^CHANGELOG.md$' && echo "Thanks for helping keep our CHANGELOG up-to-date!"
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
# ml stuff
|
||||
wandb/
|
||||
*histogram.png
|
||||
*.json
|
||||
dolma_previews/*
|
||||
s2_previews/*
|
||||
gnarly_previews/*
|
||||
s2orc_previews/*
|
||||
s2orc_previews_3200/*
|
||||
sample200_vllm/*
|
||||
sample200_sglang/*
|
||||
pdelfin_testset/*
|
||||
localworkspace/*
|
||||
math_data/*
|
||||
math_data_big/*
|
||||
gpt4otestset/*
|
||||
old_train/
|
||||
gpt4otestset_output/*
|
||||
pdfs/*
|
||||
olmOCR-bench/*
|
||||
table_data*/
|
||||
/synth*/
|
||||
dolma_samples/*
|
||||
old_train/
|
||||
filtered_items/
|
||||
filtered_items_prefilter/
|
||||
augraphy_cache/
|
||||
/*.html
|
||||
html_templates*/
|
||||
olmocr-synthmix*/
|
||||
scoreelo.csv
|
||||
debug.log
|
||||
birrpipeline-debug.log
|
||||
beakerpipeline-debug.log
|
||||
olmocr-pipeline-debug.log
|
||||
|
||||
# build artifacts
|
||||
|
||||
.eggs/
|
||||
.mypy_cache
|
||||
*.egg-info/
|
||||
build/
|
||||
dist/
|
||||
pip-wheel-metadata/
|
||||
|
||||
|
||||
# dev tools
|
||||
|
||||
.envrc
|
||||
.python-version
|
||||
.idea
|
||||
.venv/
|
||||
.vscode/
|
||||
/*.iml
|
||||
pyrightconfig.json
|
||||
|
||||
|
||||
# jupyter notebooks
|
||||
|
||||
.ipynb_checkpoints
|
||||
|
||||
|
||||
# miscellaneous
|
||||
|
||||
.cache/
|
||||
doc/_build/
|
||||
*.swp
|
||||
.DS_Store
|
||||
|
||||
|
||||
# python
|
||||
|
||||
*.pyc
|
||||
*.pyo
|
||||
__pycache__
|
||||
|
||||
|
||||
# testing and continuous integration
|
||||
|
||||
.coverage
|
||||
.pytest_cache/
|
||||
.benchmarks
|
||||
|
||||
# documentation build artifacts
|
||||
|
||||
docs/build
|
||||
site/
|
||||
@@ -0,0 +1,14 @@
|
||||
version: 2
|
||||
|
||||
sphinx:
|
||||
configuration: docs/source/conf.py
|
||||
fail_on_warning: true
|
||||
|
||||
python:
|
||||
version: "3.8"
|
||||
install:
|
||||
- method: pip
|
||||
path: .
|
||||
extra_requirements:
|
||||
- dev
|
||||
|
||||
+116
@@ -0,0 +1,116 @@
|
||||
# Changelog
|
||||
|
||||
All notable changes to this project will be documented in this file.
|
||||
|
||||
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
|
||||
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
||||
|
||||
## Unreleased
|
||||
|
||||
## [v0.4.27](https://github.com/allenai/olmocr/releases/tag/v0.4.27) - 2026-03-12
|
||||
|
||||
## [v0.4.26](https://github.com/allenai/olmocr/releases/tag/v0.4.26) - 2026-03-06
|
||||
|
||||
## [v0.4.25](https://github.com/allenai/olmocr/releases/tag/v0.4.25) - 2026-01-25
|
||||
|
||||
## [v0.4.24](https://github.com/allenai/olmocr/releases/tag/v0.4.24) - 2026-01-23
|
||||
|
||||
## [v0.4.23](https://github.com/allenai/olmocr/releases/tag/v0.4.23) - 2026-01-23
|
||||
|
||||
## [v0.4.22](https://github.com/allenai/olmocr/releases/tag/v0.4.22) - 2026-01-23
|
||||
|
||||
## [v0.4.21](https://github.com/allenai/olmocr/releases/tag/v0.4.21) - 2026-01-23
|
||||
|
||||
## [v0.4.20](https://github.com/allenai/olmocr/releases/tag/v0.4.20) - 2026-01-20
|
||||
|
||||
## [v0.4.19](https://github.com/allenai/olmocr/releases/tag/v0.4.19) - 2026-01-20
|
||||
|
||||
## [v0.4.18](https://github.com/allenai/olmocr/releases/tag/v0.4.18) - 2026-01-20
|
||||
|
||||
## [v0.4.17](https://github.com/allenai/olmocr/releases/tag/v0.4.17) - 2026-01-20
|
||||
|
||||
## [v0.4.16](https://github.com/allenai/olmocr/releases/tag/v0.4.16) - 2026-01-12
|
||||
|
||||
## [v0.4.15](https://github.com/allenai/olmocr/releases/tag/v0.4.15) - 2026-01-09
|
||||
|
||||
## [v0.4.14](https://github.com/allenai/olmocr/releases/tag/v0.4.14) - 2026-01-08
|
||||
|
||||
## [v0.4.13](https://github.com/allenai/olmocr/releases/tag/v0.4.13) - 2026-01-08
|
||||
|
||||
## [v0.4.12](https://github.com/allenai/olmocr/releases/tag/v0.4.12) - 2025-12-10
|
||||
|
||||
## [v0.4.11](https://github.com/allenai/olmocr/releases/tag/v0.4.11) - 2025-12-10
|
||||
|
||||
## [v0.4.10](https://github.com/allenai/olmocr/releases/tag/v0.4.10) - 2025-12-09
|
||||
|
||||
## [v0.4.9](https://github.com/allenai/olmocr/releases/tag/v0.4.9) - 2025-12-09
|
||||
|
||||
## [v0.4.7](https://github.com/allenai/olmocr/releases/tag/v0.4.7) - 2025-12-01
|
||||
|
||||
## [v0.4.6](https://github.com/allenai/olmocr/releases/tag/v0.4.6) - 2025-11-17
|
||||
|
||||
## [v0.4.5](https://github.com/allenai/olmocr/releases/tag/v0.4.5) - 2025-11-14
|
||||
|
||||
## [v0.4.4](https://github.com/allenai/olmocr/releases/tag/v0.4.4) - 2025-11-04
|
||||
|
||||
## [v0.4.3](https://github.com/allenai/olmocr/releases/tag/v0.4.3) - 2025-10-31
|
||||
|
||||
## [v0.4.2](https://github.com/allenai/olmocr/releases/tag/v0.4.2) - 2025-10-22
|
||||
|
||||
## [v0.4.1](https://github.com/allenai/olmocr/releases/tag/v0.4.1) - 2025-10-22
|
||||
|
||||
## [v0.4.0](https://github.com/allenai/olmocr/releases/tag/v0.4.0) - 2025-10-22
|
||||
|
||||
## [v0.3.9](https://github.com/allenai/olmocr/releases/tag/v0.3.9) - 2025-10-07
|
||||
|
||||
## [v0.3.8](https://github.com/allenai/olmocr/releases/tag/v0.3.8) - 2025-10-06
|
||||
|
||||
## [v0.3.7](https://github.com/allenai/olmocr/releases/tag/v0.3.7) - 2025-10-06
|
||||
|
||||
## [v0.3.6](https://github.com/allenai/olmocr/releases/tag/v0.3.6) - 2025-09-29
|
||||
|
||||
## [v0.3.4](https://github.com/allenai/olmocr/releases/tag/v0.3.4) - 2025-08-31
|
||||
|
||||
## [v0.3.3](https://github.com/allenai/olmocr/releases/tag/v0.3.3) - 2025-08-15
|
||||
|
||||
## [v0.3.2](https://github.com/allenai/olmocr/releases/tag/v0.3.2) - 2025-08-14
|
||||
|
||||
## [v0.3.1](https://github.com/allenai/olmocr/releases/tag/v0.3.1) - 2025-08-14
|
||||
|
||||
## [v0.3.0](https://github.com/allenai/olmocr/releases/tag/v0.3.0) - 2025-08-13
|
||||
|
||||
## [v0.2.3](https://github.com/allenai/olmocr/releases/tag/v0.2.3) - 2025-08-04
|
||||
|
||||
## [v0.2.2](https://github.com/allenai/olmocr/releases/tag/v0.2.2) - 2025-08-04
|
||||
|
||||
## [v0.2.1](https://github.com/allenai/olmocr/releases/tag/v0.2.1) - 2025-07-23
|
||||
|
||||
## [v0.2.0](https://github.com/allenai/olmocr/releases/tag/v0.2.0) - 2025-07-23
|
||||
|
||||
## [v0.1.76](https://github.com/allenai/olmocr/releases/tag/v0.1.76) - 2025-06-23
|
||||
|
||||
## [v0.1.75](https://github.com/allenai/olmocr/releases/tag/v0.1.75) - 2025-06-17
|
||||
|
||||
## [v0.1.74](https://github.com/allenai/olmocr/releases/tag/v0.1.74) - 2025-06-17
|
||||
|
||||
## [v0.1.73](https://github.com/allenai/olmocr/releases/tag/v0.1.73) - 2025-06-17
|
||||
|
||||
## [v0.1.72](https://github.com/allenai/olmocr/releases/tag/v0.1.72) - 2025-06-17
|
||||
|
||||
## [v0.1.71](https://github.com/allenai/olmocr/releases/tag/v0.1.71) - 2025-05-30
|
||||
|
||||
## [v0.1.70](https://github.com/allenai/olmocr/releases/tag/v0.1.70) - 2025-05-23
|
||||
|
||||
## [v0.1.69](https://github.com/allenai/olmocr/releases/tag/v0.1.69) - 2025-05-20
|
||||
|
||||
## [v0.1.68](https://github.com/allenai/olmocr/releases/tag/v0.1.68) - 2025-05-19
|
||||
|
||||
## [v0.1.60](https://github.com/allenai/olmocr/releases/tag/v0.1.60) - 2025-03-17
|
||||
|
||||
## [v0.1.58](https://github.com/allenai/olmocr/releases/tag/v0.1.58) - 2025-02-15
|
||||
|
||||
## [v0.1.53](https://github.com/allenai/olmocr/releases/tag/v0.1.53) - 2025-02-14
|
||||
|
||||
- Fixed git checks
|
||||
|
||||
- Added gemini and claude runners and a viewer.
|
||||
+59
@@ -0,0 +1,59 @@
|
||||
FROM vllm/vllm-openai:v0.11.2
|
||||
|
||||
ENV PYTHON_VERSION=3.12
|
||||
ENV CUSTOM_PY="/usr/bin/python${PYTHON_VERSION}"
|
||||
|
||||
# Workaround for installing fonts, which are needed for good rendering of documents
|
||||
RUN DIST_PY=$(ls /usr/bin/python3.[0-9]* | sort -V | head -n1) && \
|
||||
# If a python alternative scheme already exists, remember its value so we \
|
||||
# can restore it later; otherwise, we will restore to CUSTOM_PY when we \
|
||||
# are done. \
|
||||
if update-alternatives --query python3 >/dev/null 2>&1; then \
|
||||
ORIGINAL_PY=$(update-alternatives --query python3 | awk -F": " '/Value:/ {print $2}'); \
|
||||
else \
|
||||
ORIGINAL_PY=$CUSTOM_PY; \
|
||||
fi && \
|
||||
# ---- APT operations that require the distro python3 ------------------- \
|
||||
echo "Temporarily switching python3 alternative to ${DIST_PY} so that APT scripts use the distro‑built Python runtime." && \
|
||||
update-alternatives --install /usr/bin/python3 python3 ${DIST_PY} 1 && \
|
||||
update-alternatives --set python3 ${DIST_PY} && \
|
||||
update-alternatives --install /usr/bin/python python ${DIST_PY} 1 && \
|
||||
update-alternatives --set python ${DIST_PY} && \
|
||||
apt-get update -y && \
|
||||
apt-get remove -y python3-blinker || true && \
|
||||
# Pre‑seed the Microsoft Core Fonts EULA so the build is non‑interactive \
|
||||
echo "ttf-mscorefonts-installer msttcorefonts/accepted-mscorefonts-eula select true" | debconf-set-selections && \
|
||||
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
|
||||
python3-apt \
|
||||
update-notifier-common \
|
||||
poppler-utils \
|
||||
fonts-crosextra-caladea \
|
||||
fonts-crosextra-carlito \
|
||||
gsfonts \
|
||||
lcdf-typetools \
|
||||
ttf-mscorefonts-installer \
|
||||
git git-lfs curl wget unzip && \
|
||||
# ---- Restore the original / custom Python alternative ----------------- \
|
||||
echo "Restoring python3 alternative to ${ORIGINAL_PY}" && \
|
||||
update-alternatives --install /usr/bin/python3 python3 ${ORIGINAL_PY} 1 && \
|
||||
update-alternatives --set python3 ${ORIGINAL_PY} && \
|
||||
update-alternatives --install /usr/bin/python python ${ORIGINAL_PY} 1 || true && \
|
||||
update-alternatives --set python ${ORIGINAL_PY} || true && \
|
||||
# Ensure pip is available for the restored Python \
|
||||
curl -sS https://bootstrap.pypa.io/get-pip.py | ${ORIGINAL_PY}
|
||||
|
||||
# keep the build context clean
|
||||
WORKDIR /build
|
||||
COPY . /build
|
||||
|
||||
# Needed to resolve setuptools dependencies
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
RUN uv pip install --system --no-cache ".[bench]"
|
||||
|
||||
RUN playwright install-deps
|
||||
RUN playwright install chromium
|
||||
|
||||
RUN python3 -m olmocr.pipeline --help
|
||||
|
||||
# Override the vLLM base image's entrypoint to allow interactive bash access
|
||||
ENTRYPOINT ["/bin/bash"]
|
||||
@@ -0,0 +1,30 @@
|
||||
# Build from the base olmocr image
|
||||
FROM alleninstituteforai/olmocr:latest
|
||||
|
||||
# Allow specifying which model to include at build time
|
||||
# Default to the latest olmOCR model
|
||||
ARG MODEL_NAME=allenai/olmOCR-2-7B-1025-FP8
|
||||
|
||||
# Set model cache directory to a fixed location in the image
|
||||
ENV HF_HOME=/opt/models
|
||||
ENV TRANSFORMERS_CACHE=/opt/models
|
||||
ENV OLMOCR_MODEL=${MODEL_NAME}
|
||||
|
||||
# Pre-download the olmOCR model into the image
|
||||
# This adds ~16GB to the image size but eliminates runtime downloads
|
||||
RUN python -c "from huggingface_hub import snapshot_download; \
|
||||
import os; \
|
||||
model = os.environ.get('OLMOCR_MODEL'); \
|
||||
print(f'Downloading model: {model}'); \
|
||||
snapshot_download(model, cache_dir='/opt/models/hub'); \
|
||||
print(f'Model {model} successfully downloaded and cached in image')"
|
||||
|
||||
# Verify the model is present
|
||||
RUN python -c "import os; \
|
||||
model = os.environ.get('OLMOCR_MODEL').replace('/', '--'); \
|
||||
model_path = f'/opt/models/hub/models--{model}'; \
|
||||
assert os.path.exists(model_path), f'Model not found at {model_path}'; \
|
||||
size = sum(os.path.getsize(os.path.join(dp, f)) for dp, dn, fn in os.walk(model_path) for f in fn) / (1024**3); \
|
||||
print(f'Model size: {size:.2f} GB')"
|
||||
|
||||
# The entrypoint is already set to /bin/bash in the base image
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
https://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "{}"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright {yyyy} {name of copyright owner}
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,17 @@
|
||||
.PHONY : docs
|
||||
docs :
|
||||
rm -rf docs/build/
|
||||
sphinx-autobuild -b html --watch olmocr/ docs/source/ docs/build/
|
||||
|
||||
.PHONY : run-checks
|
||||
run-checks :
|
||||
isort --check .
|
||||
black --check .
|
||||
ruff check .
|
||||
mypy .
|
||||
CUDA_VISIBLE_DEVICES='' pytest -v --color=yes --doctest-modules tests/ olmocr/
|
||||
|
||||
.PHONY : build
|
||||
build :
|
||||
rm -rf *.egg-info/
|
||||
python -m build
|
||||
@@ -0,0 +1,551 @@
|
||||
<div align="center">
|
||||
<img width="350" alt="olmocr-2-full@2x" src="https://github.com/user-attachments/assets/24f1b596-4059-46f1-8130-5d72dcc0b02e" />
|
||||
<hr/>
|
||||
</div>
|
||||
<p align="center">
|
||||
<a href="https://github.com/allenai/OLMo/blob/main/LICENSE">
|
||||
<img alt="GitHub License" src="https://img.shields.io/github/license/allenai/OLMo">
|
||||
</a>
|
||||
<a href="https://github.com/allenai/olmocr/releases">
|
||||
<img alt="GitHub release" src="https://img.shields.io/github/release/allenai/olmocr.svg">
|
||||
</a>
|
||||
<a href="https://arxiv.org/abs/2502.18443">
|
||||
<img alt="Tech Report v1" src="https://img.shields.io/badge/Paper_v1-olmOCR-blue">
|
||||
</a>
|
||||
<a href="https://arxiv.org/abs/2510.19817">
|
||||
<img alt="Tech Report v2" src="https://img.shields.io/badge/Paper_v2-olmOCR-blue">
|
||||
</a>
|
||||
<a href="https://olmocr.allenai.org">
|
||||
<img alt="Demo" src="https://img.shields.io/badge/Ai2-Demo-F0529C">
|
||||
</a>
|
||||
<a href="https://discord.gg/sZq3jTNVNG">
|
||||
<img alt="Discord" src="https://img.shields.io/badge/Discord%20-%20blue?style=flat&logo=discord&label=Ai2&color=%235B65E9">
|
||||
</a>
|
||||
</p>
|
||||
|
||||
A toolkit for converting PDFs and other image-based document formats into clean, readable, plain text format.
|
||||
|
||||
Try the online demo: [https://olmocr.allenai.org/](https://olmocr.allenai.org/)
|
||||
|
||||
Features:
|
||||
- Convert PDF, PNG, and JPEG based documents into clean Markdown
|
||||
- Support for equations, tables, handwriting, and complex formatting
|
||||
- Automatically removes headers and footers
|
||||
- Convert into text with a natural reading order, even in the presence of
|
||||
figures, multi-column layouts, and insets
|
||||
- Efficient, less than $200 USD per million pages converted
|
||||
- (Based on a 7B parameter VLM, so it requires a GPU)
|
||||
|
||||
### News
|
||||
- October 21, 2025 - v0.4.0 - [New model release](https://huggingface.co/allenai/olmOCR-2-7B-1025-FP8), boosts olmOCR-bench score by ~4 points using synthetic data and introduces RL training.
|
||||
- August 13, 2025 - v0.3.0 - [New model release](https://huggingface.co/allenai/olmOCR-7B-0825-FP8), fixes auto-rotation detection, and hallucinations on blank documents.
|
||||
- July 24, 2025 - v0.2.1 - [New model release](https://huggingface.co/allenai/olmOCR-7B-0725-FP8), scores 3 points higher on [olmOCR-Bench](https://github.com/allenai/olmocr/tree/main/olmocr/bench), also runs significantly faster because it's default FP8, and needs much fewer retries per document.
|
||||
- July 23, 2025 - v0.2.0 - New cleaned up [trainer code](https://github.com/allenai/olmocr/tree/main/olmocr/train), makes it much simpler to train olmOCR models yourself.
|
||||
- June 17, 2025 - v0.1.75 - Switch from sglang to vllm based inference pipeline, updated docker image to CUDA 12.8.
|
||||
- May 23, 2025 - v0.1.70 - Official docker support and images are now available! [See Docker usage](#using-docker)
|
||||
- May 19, 2025 - v0.1.68 - [olmOCR-Bench](https://github.com/allenai/olmocr/tree/main/olmocr/bench) launch, scoring 77.4. Launch includes 2 point performance boost in olmOCR pipeline due to bug fixes with prompts.
|
||||
- Mar 17, 2025 - v0.1.60 - Performance improvements due to better temperature selection in sampling.
|
||||
- Feb 25, 2025 - v0.1.58 - Initial public launch and demo.
|
||||
|
||||
### Benchmark
|
||||
|
||||
[**olmOCR-Bench**](https://github.com/allenai/olmocr/tree/main/olmocr/bench):
|
||||
We also ship a comprehensive benchmark suite covering over 7,000 test cases across 1,400 documents to help measure performance of OCR systems.
|
||||
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th></th>
|
||||
<th>ArXiv</th>
|
||||
<th>Old<br>scans<br>math</th>
|
||||
<th>Tables</th>
|
||||
<th>Old<br>scans</th>
|
||||
<th>Headers<br>&<br>footers</th>
|
||||
<th>Multi<br>column</th>
|
||||
<th>Long<br>tiny<br>text</th>
|
||||
<th>Base</th>
|
||||
<th>Overall</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>Mistral OCR API</td>
|
||||
<td>77.2</td>
|
||||
<td>67.5</td>
|
||||
<td>60.6</td>
|
||||
<td>29.3</td>
|
||||
<td>93.6</td>
|
||||
<td>71.3</td>
|
||||
<td>77.1</td>
|
||||
<td>99.4</td>
|
||||
<td>72.0±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Marker 1.10.1</td>
|
||||
<td>83.8</td>
|
||||
<td>66.8</td>
|
||||
<td>72.9</td>
|
||||
<td>33.5</td>
|
||||
<td>86.6</td>
|
||||
<td>80.0</td>
|
||||
<td>85.7</td>
|
||||
<td>99.3</td>
|
||||
<td>76.1±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>MinerU 2.5.4*</td>
|
||||
<td>76.6</td>
|
||||
<td>54.6</td>
|
||||
<td>84.9</td>
|
||||
<td>33.7</td>
|
||||
<td>96.6</td>
|
||||
<td>78.2</td>
|
||||
<td>83.5</td>
|
||||
<td>93.7</td>
|
||||
<td>75.2±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>DeepSeek-OCR</td>
|
||||
<td>77.2</td>
|
||||
<td>73.6</td>
|
||||
<td>80.2</td>
|
||||
<td>33.3</td>
|
||||
<td>96.1</td>
|
||||
<td>66.4</td>
|
||||
<td>79.4</td>
|
||||
<td>99.8</td>
|
||||
<td>75.7±1.0</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Nanonets-OCR2-3B</td>
|
||||
<td>75.4</td>
|
||||
<td>46.1</td>
|
||||
<td>86.8</td>
|
||||
<td>40.9</td>
|
||||
<td>32.1</td>
|
||||
<td>81.9</td>
|
||||
<td>93.0</td>
|
||||
<td>99.6</td>
|
||||
<td>69.5±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>PaddleOCR-VL*</td>
|
||||
<td>85.7</td>
|
||||
<td>71.0</td>
|
||||
<td>84.1</td>
|
||||
<td>37.8</td>
|
||||
<td>97.0</td>
|
||||
<td>79.9</td>
|
||||
<td>85.7</td>
|
||||
<td>98.5</td>
|
||||
<td>80.0±1.0</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Infinity-Parser 7B*</td>
|
||||
<td>84.4</td>
|
||||
<td>83.8</td>
|
||||
<td>85.0</td>
|
||||
<td>47.9</td>
|
||||
<td>88.7</td>
|
||||
<td>84.2</td>
|
||||
<td>86.4</td>
|
||||
<td>99.8</td>
|
||||
<td>82.5±?</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Chandra OCR 0.1.0*</td>
|
||||
<td>82.2</td>
|
||||
<td>80.3</td>
|
||||
<td>88.0</td>
|
||||
<td>50.4</td>
|
||||
<td>90.8</td>
|
||||
<td>81.2</td>
|
||||
<td>92.3</td>
|
||||
<td>99.9</td>
|
||||
<td>83.1±0.9</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td colspan="10"><hr></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><strong>olmOCR v0.4.0</strong></td>
|
||||
<td>83.0</td>
|
||||
<td>82.3</td>
|
||||
<td>84.9</td>
|
||||
<td>47.7</td>
|
||||
<td>96.1</td>
|
||||
<td>83.7</td>
|
||||
<td>81.9</td>
|
||||
<td>99.7</td>
|
||||
<td>82.4±1.1</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
|
||||
### Installation
|
||||
|
||||
#### System Dependencies
|
||||
|
||||
You will need to install poppler-utils and additional fonts for rendering PDF images.
|
||||
|
||||
Install dependencies (Ubuntu/Debian):
|
||||
```bash
|
||||
sudo apt-get update
|
||||
sudo apt-get install poppler-utils ttf-mscorefonts-installer msttcorefonts fonts-crosextra-caladea fonts-crosextra-carlito gsfonts lcdf-typetools
|
||||
```
|
||||
|
||||
#### Python Installation
|
||||
|
||||
Set up a conda environment and install olmocr. The requirements for running olmOCR
|
||||
are difficult to install in an existing python environment, so please do make a clean python environment to install into.
|
||||
|
||||
```bash
|
||||
conda create -n olmocr python=3.11
|
||||
conda activate olmocr
|
||||
```
|
||||
|
||||
Choose the installation option that matches your use case:
|
||||
|
||||
**Option 1: Remote Inference (Lightweight)**
|
||||
|
||||
If you plan to use a remote vLLM server with the `--server` flag, install the base package:
|
||||
```bash
|
||||
pip install olmocr
|
||||
```
|
||||
This avoids installing heavy GPU dependencies like PyTorch (~2GB+).
|
||||
|
||||
**Option 2: Local GPU Inference**
|
||||
|
||||
Requirements:
|
||||
- Recent NVIDIA GPU (tested on RTX 4090, L40S, A100, H100) with at least 12 GB of GPU RAM
|
||||
- 30GB of free disk space
|
||||
|
||||
For running inference with your own GPU:
|
||||
```bash
|
||||
pip install olmocr[gpu] --extra-index-url https://download.pytorch.org/whl/cu128
|
||||
|
||||
# Recommended: Install flash infer for faster inference on GPU
|
||||
pip install https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.5%2Bcu128torch2.7-cp38-abi3-linux_x86_64.whl
|
||||
```
|
||||
|
||||
**Option 3: Beaker Cluster Execution**
|
||||
|
||||
For submitting jobs to Beaker clusters with the `--beaker` flag:
|
||||
```bash
|
||||
pip install olmocr[beaker]
|
||||
```
|
||||
|
||||
**Option 4: Benchmark Suite**
|
||||
|
||||
For running the olmOCR benchmark suite:
|
||||
```bash
|
||||
pip install olmocr[bench]
|
||||
```
|
||||
|
||||
**Combined Installation**
|
||||
|
||||
You can combine multiple options:
|
||||
```bash
|
||||
# GPU + Beaker support
|
||||
pip install olmocr[gpu,beaker] --extra-index-url https://download.pytorch.org/whl/cu128
|
||||
|
||||
# GPU + Benchmark support
|
||||
pip install olmocr[gpu,bench] --extra-index-url https://download.pytorch.org/whl/cu128
|
||||
```
|
||||
|
||||
**Troubleshooting**
|
||||
|
||||
If you run into errors about `too many open files`, update your ulimit:
|
||||
```bash
|
||||
ulimit -n 65536
|
||||
```
|
||||
|
||||
### Usage Examples
|
||||
|
||||
For quick testing, try the [web demo](https://olmocr.allen.ai/).
|
||||
|
||||
**Convert a Single PDF (Local GPU):**
|
||||
```bash
|
||||
# Download a sample PDF
|
||||
curl -o olmocr-sample.pdf https://olmocr.allenai.org/papers/olmocr_3pg_sample.pdf
|
||||
|
||||
# Convert it to markdown
|
||||
olmocr ./localworkspace --markdown --pdfs olmocr-sample.pdf
|
||||
```
|
||||
|
||||
**Convert an Image file:**
|
||||
```bash
|
||||
olmocr ./localworkspace --markdown --pdfs random_page.png
|
||||
```
|
||||
|
||||
**Convert Multiple PDFs:**
|
||||
```bash
|
||||
olmocr ./localworkspace --markdown --pdfs tests/gnarly_pdfs/*.pdf
|
||||
```
|
||||
|
||||
**Use Remote Inference Server:**
|
||||
```bash
|
||||
olmocr ./localworkspace --server http://remote-server:8000/v1 --model allenai/olmOCR-2-7B-1025-FP8 --markdown --pdfs *.pdf
|
||||
```
|
||||
|
||||
With the `--markdown` flag, results will be stored as markdown files inside of `./localworkspace/markdown/`.
|
||||
|
||||
> **Note:** You can also use `python -m olmocr.pipeline` instead of `olmocr` if you prefer.
|
||||
|
||||
#### Viewing Results
|
||||
|
||||
The `./localworkspace/` workspace folder will then have both [Dolma](https://github.com/allenai/dolma) and markdown files (if using `--markdown`).
|
||||
|
||||
|
||||
```bash
|
||||
cat localworkspace/markdown/olmocr-sample.md
|
||||
```
|
||||
|
||||
```
|
||||
olmOCR: Unlocking Trillions of Tokens in PDFs with Vision Language Models
|
||||
...
|
||||
```
|
||||
|
||||
### Using an Inference Provider or External Server
|
||||
|
||||
If you have a vLLM server already running elsewhere (or any inference platform implementing the OpenAI API), you can point olmOCR to use it instead of spawning a local instance.
|
||||
|
||||
**Installation for Remote Inference:**
|
||||
```bash
|
||||
# Lightweight installation - no GPU dependencies needed
|
||||
pip install olmocr
|
||||
```
|
||||
|
||||
**Using an External Server:**
|
||||
```bash
|
||||
# Use external vLLM server instead of local one
|
||||
olmocr ./localworkspace --server http://remote-server:8000/v1 --model allenai/olmOCR-2-7B-1025-FP8 --markdown --pdfs tests/gnarly_pdfs/*.pdf
|
||||
```
|
||||
|
||||
The served model name in vLLM needs to match the value provided in `--model`.
|
||||
|
||||
**Example vLLM Server Launch:**
|
||||
```bash
|
||||
vllm serve allenai/olmOCR-2-7B-1025-FP8 --max-model-len 16384
|
||||
```
|
||||
|
||||
#### Verified External Providers
|
||||
|
||||
We have tested `olmOCR-2-7B-1025-FP8` on these external model providers and confirmed that they work
|
||||
|
||||
| | $/1M Input tokens | $/1M Output tokens | Example Command |
|
||||
|-----------------------------------------------------------------------------|-------------------|--------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| [Cirrascale](https://ai2endpoints.cirrascale.ai/models/overview) | $0.07 | $0.15 | `olmocr ./workspace --server https://ai2endpoints.cirrascale.ai/api --api_key sk-XXXXXXX --workers 1 --max_concurrent_requests 20 --model olmOCR-2-7B-1025 --pdfs tests/gnarly_pdfs/*.pdf` |
|
||||
| [DeepInfra](https://deepinfra.com/) | $0.09 | $0.19 | `olmocr ./workspace --server https://api.deepinfra.com/v1/openai --api_key DfXXXXXXX --workers 1 --max_concurrent_requests 20 --model allenai/olmOCR-2-7B-1025 --pdfs tests/gnarly_pdfs/*.pdf` |
|
||||
| [Parasail](https://www.saas.parasail.io/serverless?name=olmocr-7b-1025-fp8) | $0.10 | $0.20 | `olmocr ./workspace --server https://api.parasail.io/v1 --api_key psk-XXXXX --workers 1 --max_concurrent_requests 20 --model allenai/olmOCR-2-7B-1025 --pdfs tests/gnarly_pdfs/*.pdf` |
|
||||
|
||||
|
||||
Notes on arguments
|
||||
- `--server`: Defines the OpenAI-compatible endpoint: ex `https://api.deepinfra.com/v1/openai`
|
||||
- `--api_key`: Your API key, bassed in via Authorization Bearer HTTP header
|
||||
- `--max_concurrent_requests`: Max concurrent requests that will be in-flight to the inference provider at one time
|
||||
- `--workers`: Max number of page groups that will be processed at once. You may want to set this to `1` so that you finish one group of stuff before moving on.
|
||||
- `--pages_per_group`: You may want a smaller number of pages per group as many external provides have lower concurrent request limits
|
||||
- `--model`: The model identifier, ex. `allenai/olmOCR-2-7B-1025`, different providers have different names, and if you run locally, you can use `olmocr`
|
||||
- Other arguments work the same as with local inference
|
||||
|
||||
|
||||
### Multi-node / Cluster Usage
|
||||
|
||||
If you want to convert millions of PDFs using multiple nodes running in parallel, olmOCR supports
|
||||
reading PDFs from AWS S3 and coordinating work using an AWS S3 output bucket.
|
||||
|
||||
**Start the first worker node:**
|
||||
```bash
|
||||
olmocr s3://my_s3_bucket/pdfworkspaces/exampleworkspace --pdfs s3://my_s3_bucket/jakep/gnarly_pdfs/*.pdf
|
||||
```
|
||||
|
||||
This sets up a simple work queue in your AWS bucket and starts converting PDFs.
|
||||
|
||||
**On subsequent worker nodes:**
|
||||
```bash
|
||||
olmocr s3://my_s3_bucket/pdfworkspaces/exampleworkspace
|
||||
```
|
||||
|
||||
They will automatically start grabbing items from the same workspace queue.
|
||||
|
||||
#### Using Beaker for Cluster Execution
|
||||
|
||||
If you are at Ai2 and want to linearize millions of PDFs efficiently using [beaker](https://www.beaker.org), install with Beaker support:
|
||||
|
||||
```bash
|
||||
pip install olmocr[gpu,beaker] --extra-index-url https://download.pytorch.org/whl/cu128
|
||||
```
|
||||
|
||||
Then use the `--beaker` flag to prepare the workspace locally and launch N GPU workers in the cluster:
|
||||
|
||||
```bash
|
||||
olmocr s3://my_s3_bucket/pdfworkspaces/exampleworkspace --pdfs s3://my_s3_bucket/jakep/gnarly_pdfs/*.pdf --beaker --beaker_gpus 4
|
||||
```
|
||||
|
||||
|
||||
### Using Docker
|
||||
|
||||
Pull the Docker image (large, includes the model, ~30GB):
|
||||
```bash
|
||||
docker pull alleninstituteforai/olmocr:latest-with-model
|
||||
```
|
||||
|
||||
For advanced users who want to manage their own model downloads, we also provide a base image without the model:
|
||||
```bash
|
||||
docker pull alleninstituteforai/olmocr:latest
|
||||
```
|
||||
|
||||
#### Quick Start - Process PDFs
|
||||
|
||||
Process a single PDF in your current directory:
|
||||
```bash
|
||||
docker run --gpus all \
|
||||
-v $(pwd):/workspace \
|
||||
alleninstituteforai/olmocr:latest-with-model \
|
||||
-c "olmocr /workspace/output --markdown --pdfs /workspace/sample.pdf"
|
||||
```
|
||||
|
||||
Process multiple PDFs:
|
||||
```bash
|
||||
docker run --gpus all \
|
||||
-v /path/to/pdfs:/input \
|
||||
-v /path/to/output:/output \
|
||||
alleninstituteforai/olmocr:latest-with-model \
|
||||
-c "olmocr /output --markdown --pdfs /input/*.pdf"
|
||||
```
|
||||
|
||||
#### Interactive Mode
|
||||
|
||||
Run the container interactively for exploration and debugging:
|
||||
```bash
|
||||
docker run -it --gpus all alleninstituteforai/olmocr:latest-with-model
|
||||
```
|
||||
|
||||
> Visit our Docker repository on [Docker Hub](https://hub.docker.com/r/alleninstituteforai/olmocr) for more information.
|
||||
|
||||
### Full Documentation
|
||||
|
||||
To see all available options:
|
||||
```bash
|
||||
olmocr --help
|
||||
usage: pipeline.py [-h] [--pdfs [PDFS ...]] [--model MODEL] [--workspace_profile WORKSPACE_PROFILE] [--pdf_profile PDF_PROFILE] [--pages_per_group PAGES_PER_GROUP] [--max_page_retries MAX_PAGE_RETRIES] [--max_page_error_rate MAX_PAGE_ERROR_RATE] [--workers WORKERS]
|
||||
[--apply_filter] [--stats] [--markdown] [--target_longest_image_dim TARGET_LONGEST_IMAGE_DIM] [--target_anchor_text_len TARGET_ANCHOR_TEXT_LEN] [--guided_decoding] [--gpu-memory-utilization GPU_MEMORY_UTILIZATION] [--max_model_len MAX_MODEL_LEN]
|
||||
[--tensor-parallel-size TENSOR_PARALLEL_SIZE] [--data-parallel-size DATA_PARALLEL_SIZE] [--port PORT] [--server SERVER] [--beaker] [--beaker_workspace BEAKER_WORKSPACE] [--beaker_cluster BEAKER_CLUSTER] [--beaker_gpus BEAKER_GPUS] [--beaker_priority BEAKER_PRIORITY]
|
||||
workspace
|
||||
|
||||
Manager for running millions of PDFs through a batch inference pipeline
|
||||
|
||||
positional arguments:
|
||||
workspace The filesystem path where work will be stored, can be a local folder, or an s3 path if coordinating work with many workers, s3://bucket/prefix/
|
||||
|
||||
options:
|
||||
-h, --help show this help message and exit
|
||||
--pdfs [PDFS ...] Path to add pdfs stored in s3 to the workspace, can be a glob path s3://bucket/prefix/*.pdf or path to file containing list of pdf paths
|
||||
--model MODEL Path where the model is located, allenai/olmOCR-7B-0725-FP8 is the default, can be local, s3, or hugging face.
|
||||
--workspace_profile WORKSPACE_PROFILE
|
||||
S3 configuration profile for accessing the workspace
|
||||
--pdf_profile PDF_PROFILE
|
||||
S3 configuration profile for accessing the raw pdf documents
|
||||
--pages_per_group PAGES_PER_GROUP
|
||||
Aiming for this many pdf pages per work item group
|
||||
--max_page_retries MAX_PAGE_RETRIES
|
||||
Max number of times we will retry rendering a page
|
||||
--max_page_error_rate MAX_PAGE_ERROR_RATE
|
||||
Rate of allowable failed pages in a document, 1/250 by default
|
||||
--workers WORKERS Number of workers to run at a time
|
||||
--apply_filter Apply basic filtering to English pdfs which are not forms, and not likely seo spam
|
||||
--stats Instead of running any job, reports some statistics about the current workspace
|
||||
--markdown Also write natural text to markdown files preserving the folder structure of the input pdfs
|
||||
--target_longest_image_dim TARGET_LONGEST_IMAGE_DIM
|
||||
Dimension on longest side to use for rendering the pdf pages
|
||||
--target_anchor_text_len TARGET_ANCHOR_TEXT_LEN
|
||||
Maximum amount of anchor text to use (characters), not used for new models
|
||||
--guided_decoding Enable guided decoding for model YAML type outputs
|
||||
|
||||
VLLM arguments:
|
||||
--gpu-memory-utilization GPU_MEMORY_UTILIZATION
|
||||
Fraction of VRAM vLLM may pre-allocate for KV-cache (passed through to vllm serve).
|
||||
--max_model_len MAX_MODEL_LEN
|
||||
Upper bound (tokens) vLLM will allocate KV-cache for, lower if VLLM won't start
|
||||
--tensor-parallel-size TENSOR_PARALLEL_SIZE, -tp TENSOR_PARALLEL_SIZE
|
||||
Tensor parallel size for vLLM
|
||||
--data-parallel-size DATA_PARALLEL_SIZE, -dp DATA_PARALLEL_SIZE
|
||||
Data parallel size for vLLM
|
||||
--port PORT Port to use for the VLLM server
|
||||
--server SERVER URL of external vLLM (or other compatible provider)
|
||||
server (e.g., http://hostname:port). If provided,
|
||||
skips spawning local vLLM instance
|
||||
|
||||
beaker/cluster execution:
|
||||
--beaker Submit this job to beaker instead of running locally
|
||||
--beaker_workspace BEAKER_WORKSPACE
|
||||
Beaker workspace to submit to
|
||||
--beaker_cluster BEAKER_CLUSTER
|
||||
Beaker clusters you want to run on
|
||||
--beaker_gpus BEAKER_GPUS
|
||||
Number of gpu replicas to run
|
||||
--beaker_priority BEAKER_PRIORITY
|
||||
Beaker priority level for the job
|
||||
```
|
||||
|
||||
## Code overview
|
||||
|
||||
There are some nice reusable pieces of the code that may be useful for your own projects:
|
||||
- A prompting strategy to get really good natural text parsing using ChatGPT 4o - [buildsilver.py](https://github.com/allenai/olmocr/blob/main/olmocr/data/buildsilver.py)
|
||||
- Basic filtering by language and SEO spam removal - [filter.py](https://github.com/allenai/olmocr/blob/main/olmocr/filter/filter.py)
|
||||
- SFT Finetuning code for Qwen2.5-VL - [train.py](https://github.com/allenai/olmocr/blob/main/olmocr/train/train.py)
|
||||
- GRPO RL Trainer - [grpo_train.py](https://github.com/allenai/olmocr/blob/main/olmocr/train/grpo_train.py)
|
||||
- Synthetic data generation - [mine_html_templates.py](https://github.com/allenai/olmocr/blob/main/olmocr/synth/mine_html_templates.py)
|
||||
- Processing millions of PDFs through a finetuned model using VLLM - [pipeline.py](https://github.com/allenai/olmocr/blob/main/olmocr/pipeline.py)
|
||||
- Viewing [Dolma docs](https://github.com/allenai/dolma) created from PDFs - [dolmaviewer.py](https://github.com/allenai/olmocr/blob/main/olmocr/viewer/dolmaviewer.py)
|
||||
|
||||
|
||||
|
||||
## Team
|
||||
|
||||
<!-- start team -->
|
||||
|
||||
**olmOCR** is developed and maintained by the AllenNLP team, backed by [the Allen Institute for Artificial Intelligence (AI2)](https://allenai.org/).
|
||||
AI2 is a non-profit institute with the mission to contribute to humanity through high-impact AI research and engineering.
|
||||
To learn more about who specifically contributed to this codebase, see [our contributors](https://github.com/allenai/olmocr/graphs/contributors) page.
|
||||
|
||||
<!-- end team -->
|
||||
|
||||
## License
|
||||
|
||||
<!-- start license -->
|
||||
|
||||
**olmOCR** is licensed under [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0).
|
||||
A full copy of the license can be found [on GitHub](https://github.com/allenai/olmocr/blob/main/LICENSE).
|
||||
|
||||
<!-- end license -->
|
||||
|
||||
## Citing
|
||||
|
||||
For olmOCR v1 and OlmOCR-bench:
|
||||
```bibtex
|
||||
@misc{olmocrbench,
|
||||
title={{olmOCR: Unlocking Trillions of Tokens in PDFs with Vision Language Models}},
|
||||
author={Jake Poznanski and Jon Borchardt and Jason Dunkelberger and Regan Huff and Daniel Lin and Aman Rangapur and Christopher Wilhelm and Kyle Lo and Luca Soldaini},
|
||||
year={2025},
|
||||
eprint={2502.18443},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL},
|
||||
url={https://arxiv.org/abs/2502.18443},
|
||||
}
|
||||
```
|
||||
|
||||
For olmOCR v2 Unit Testing Rewards with RL:
|
||||
```bibtex
|
||||
@misc{olmocr2,
|
||||
title={olmOCR 2: Unit Test Rewards for Document OCR},
|
||||
author={Jake Poznanski and Luca Soldaini and Kyle Lo},
|
||||
year={2025},
|
||||
eprint={2510.19817},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CV},
|
||||
url={https://arxiv.org/abs/2510.19817},
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`allenai/olmocr`
|
||||
- 原始仓库：https://github.com/allenai/olmocr
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
@@ -0,0 +1,24 @@
|
||||
# GitHub Release Process
|
||||
|
||||
## Steps
|
||||
|
||||
1. Update the version in `olmocr/version.py`.
|
||||
|
||||
3. Run the release script:
|
||||
|
||||
```bash
|
||||
./scripts/release.sh
|
||||
```
|
||||
|
||||
This will commit the changes to the CHANGELOG and `version.py` files and then create a new tag in git
|
||||
which will trigger a workflow on GitHub Actions that handles the rest.
|
||||
|
||||
## Fixing a failed release
|
||||
|
||||
If for some reason the GitHub Actions release workflow failed with an error that needs to be fixed, you'll have to delete both the tag and corresponding release from GitHub. After you've pushed a fix, delete the tag from your local clone with
|
||||
|
||||
```bash
|
||||
git tag -l | xargs git tag -d && git fetch -t
|
||||
```
|
||||
|
||||
Then repeat the steps above.
|
||||
@@ -0,0 +1 @@
|
||||
build
|
||||
@@ -0,0 +1,20 @@
|
||||
# Minimal makefile for Sphinx documentation
|
||||
#
|
||||
|
||||
# You can set these variables from the command line, and also
|
||||
# from the environment for the first two.
|
||||
SPHINXOPTS ?= -W
|
||||
SPHINXBUILD ?= sphinx-build
|
||||
SOURCEDIR = source
|
||||
BUILDDIR = build
|
||||
|
||||
# Put it first so that "make" without argument is like "make help".
|
||||
help:
|
||||
@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
|
||||
|
||||
.PHONY: help Makefile
|
||||
|
||||
# Catch-all target: route all unknown targets to Sphinx using the new
|
||||
# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
|
||||
%: Makefile
|
||||
@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
|
||||
@@ -0,0 +1,35 @@
|
||||
@ECHO OFF
|
||||
|
||||
pushd %~dp0
|
||||
|
||||
REM Command file for Sphinx documentation
|
||||
|
||||
if "%SPHINXBUILD%" == "" (
|
||||
set SPHINXBUILD=sphinx-build
|
||||
)
|
||||
set SOURCEDIR=source
|
||||
set BUILDDIR=build
|
||||
|
||||
if "%1" == "" goto help
|
||||
|
||||
%SPHINXBUILD% >NUL 2>NUL
|
||||
if errorlevel 9009 (
|
||||
echo.
|
||||
echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
|
||||
echo.installed, then set the SPHINXBUILD environment variable to point
|
||||
echo.to the full path of the 'sphinx-build' executable. Alternatively you
|
||||
echo.may add the Sphinx directory to PATH.
|
||||
echo.
|
||||
echo.If you don't have Sphinx installed, grab it from
|
||||
echo.https://www.sphinx-doc.org/
|
||||
exit /b 1
|
||||
)
|
||||
|
||||
%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
|
||||
goto end
|
||||
|
||||
:help
|
||||
%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
|
||||
|
||||
:end
|
||||
popd
|
||||
Symlink
+1
@@ -0,0 +1 @@
|
||||
../../CHANGELOG.md
|
||||
Symlink
+1
@@ -0,0 +1 @@
|
||||
../../.github/CONTRIBUTING.md
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 15 KiB |
@@ -0,0 +1,121 @@
|
||||
# Configuration file for the Sphinx documentation builder.
|
||||
#
|
||||
# This file only contains a selection of the most common options. For a full
|
||||
# list see the documentation:
|
||||
# https://www.sphinx-doc.org/en/master/usage/configuration.html
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
|
||||
# -- Path setup --------------------------------------------------------------
|
||||
|
||||
# If extensions (or modules to document with autodoc) are in another directory,
|
||||
# add these directories to sys.path here. If the directory is relative to the
|
||||
# documentation root, use os.path.abspath to make it absolute, like shown here.
|
||||
#
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../"))
|
||||
|
||||
from olmocr import VERSION, VERSION_SHORT # noqa: E402
|
||||
|
||||
# -- Project information -----------------------------------------------------
|
||||
|
||||
project = "olmocr"
|
||||
copyright = f"{datetime.today().year}, Allen Institute for Artificial Intelligence"
|
||||
author = "Allen Institute for Artificial Intelligence"
|
||||
version = VERSION_SHORT
|
||||
release = VERSION
|
||||
|
||||
|
||||
# -- General configuration ---------------------------------------------------
|
||||
|
||||
# Add any Sphinx extension module names here, as strings. They can be
|
||||
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
|
||||
# ones.
|
||||
extensions = [
|
||||
"sphinx.ext.autodoc",
|
||||
"sphinx.ext.napoleon",
|
||||
"myst_parser",
|
||||
"sphinx.ext.intersphinx",
|
||||
"sphinx.ext.viewcode",
|
||||
"sphinx.ext.doctest",
|
||||
"sphinx_copybutton",
|
||||
"sphinx_autodoc_typehints",
|
||||
]
|
||||
|
||||
# Tell myst-parser to assign header anchors for h1-h3.
|
||||
myst_heading_anchors = 3
|
||||
|
||||
suppress_warnings = ["myst.header"]
|
||||
|
||||
# Add any paths that contain templates here, relative to this directory.
|
||||
templates_path = ["_templates"]
|
||||
|
||||
# List of patterns, relative to source directory, that match files and
|
||||
# directories to ignore when looking for source files.
|
||||
# This pattern also affects html_static_path and html_extra_path.
|
||||
exclude_patterns = ["_build"]
|
||||
|
||||
source_suffix = [".rst", ".md"]
|
||||
|
||||
intersphinx_mapping = {
|
||||
"python": ("https://docs.python.org/3", None),
|
||||
# Uncomment these if you use them in your codebase:
|
||||
# "torch": ("https://pytorch.org/docs/stable", None),
|
||||
# "datasets": ("https://huggingface.co/docs/datasets/master/en", None),
|
||||
# "transformers": ("https://huggingface.co/docs/transformers/master/en", None),
|
||||
}
|
||||
|
||||
# By default, sort documented members by type within classes and modules.
|
||||
autodoc_member_order = "groupwise"
|
||||
|
||||
# Include default values when documenting parameter types.
|
||||
typehints_defaults = "comma"
|
||||
|
||||
|
||||
# -- Options for HTML output -------------------------------------------------
|
||||
|
||||
# The theme to use for HTML and HTML Help pages. See the documentation for
|
||||
# a list of builtin themes.
|
||||
#
|
||||
html_theme = "furo"
|
||||
|
||||
html_title = f"olmocr v{VERSION}"
|
||||
|
||||
# Add any paths that contain custom static files (such as style sheets) here,
|
||||
# relative to this directory. They are copied after the builtin static files,
|
||||
# so a file named "default.css" will overwrite the builtin "default.css".
|
||||
html_static_path = ["_static"]
|
||||
|
||||
html_css_files = ["css/custom.css"]
|
||||
|
||||
html_favicon = "_static/favicon.ico"
|
||||
|
||||
html_theme_options = {
|
||||
"footer_icons": [
|
||||
{
|
||||
"name": "GitHub",
|
||||
"url": "https://github.com/allenai/olmocr",
|
||||
"html": """
|
||||
<svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 16 16">
|
||||
<path fill-rule="evenodd" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0 0 16 8c0-4.42-3.58-8-8-8z"></path>
|
||||
</svg>
|
||||
""", # noqa: E501
|
||||
"class": "",
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
# -- Hack to get rid of stupid warnings from sphinx_autodoc_typehints --------
|
||||
|
||||
|
||||
class ShutupSphinxAutodocTypehintsFilter(logging.Filter):
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
if "Cannot resolve forward reference" in record.msg:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
logging.getLogger("sphinx.sphinx_autodoc_typehints").addFilter(ShutupSphinxAutodocTypehintsFilter())
|
||||
@@ -0,0 +1,27 @@
|
||||
# **olmocr**
|
||||
|
||||
```{toctree}
|
||||
:maxdepth: 2
|
||||
:hidden:
|
||||
:caption: Getting started
|
||||
|
||||
installation
|
||||
overview
|
||||
```
|
||||
|
||||
```{toctree}
|
||||
:hidden:
|
||||
:caption: Development
|
||||
|
||||
CHANGELOG
|
||||
CONTRIBUTING
|
||||
License <https://raw.githubusercontent.com/allenai/olmocr/main/LICENSE>
|
||||
GitHub Repository <https://github.com/allenai/olmocr>
|
||||
```
|
||||
|
||||
## Indices and tables
|
||||
|
||||
```{eval-rst}
|
||||
* :ref:`genindex`
|
||||
* :ref:`modindex`
|
||||
```
|
||||
@@ -0,0 +1,27 @@
|
||||
Installation
|
||||
============
|
||||
|
||||
**olmocr** supports Python >= 3.8.
|
||||
|
||||
## Installing with `pip`
|
||||
|
||||
**olmocr** is available [on PyPI](https://pypi.org/project/olmocr/). Just run
|
||||
|
||||
```bash
|
||||
pip install olmocr
|
||||
```
|
||||
|
||||
## Installing from source
|
||||
|
||||
To install **olmocr** from source, first clone [the repository](https://github.com/allenai/olmocr):
|
||||
|
||||
```bash
|
||||
git clone https://github.com/allenai/olmocr.git
|
||||
cd olmocr
|
||||
```
|
||||
|
||||
Then run
|
||||
|
||||
```bash
|
||||
pip install -e .
|
||||
```
|
||||
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 273 KiB |
@@ -0,0 +1,3 @@
|
||||
Overview
|
||||
========
|
||||
|
||||
Binary file not shown.
@@ -0,0 +1 @@
|
||||
from .version import VERSION, VERSION_SHORT
|
||||
@@ -0,0 +1,322 @@
|
||||
# olmOCR-Bench
|
||||
|
||||
Dataset Link: https://huggingface.co/datasets/allenai/olmOCR-bench
|
||||
|
||||
We develop olmOCR-Bench in order to automatically and effectively evaluate document-level OCR of various tools.
|
||||
|
||||
olmOCR-Bench works by testing various "facts" about document pages at the PDF-level.
|
||||
Our intention is that each "fact" is very simple, unambiguous, and machine-checkable, similar to a unit test. For example, once your document has been OCRed, we may check that a particular sentence appears exactly somewhere on the page.
|
||||
|
||||
We stay away from soft metrics like edit distance comparisons, because they may assign lower scores for parses of the document that differ from the reference, but may in fact still be correct. For example, on a document containing multiple distinct articles: you want the text of each article to be grouped together, but the relative order of the two articles may not be critical. Also, some documents may have critical details, like switching x and y in an equation that can make all the difference in understanding, but would appear as just a single character edit in an edit-distance metric.
|
||||
|
||||
olmOCR-bench operates on single page PDFs directly. We make this choice because PDFs do preserve some digital metadata and information which may be helpful to some OCR systems. Almost any other format can be converted to a PDF, but not the reverse, so we try to preserve these original documents where possible.
|
||||
|
||||
We have run the benchmark against some contemporary OCR pipelines, but it is really easy
|
||||
to run it against your own OCR tools. Your tool just needs to support Markdown or plain text output.
|
||||
|
||||
<div align="center">
|
||||
<img src="https://github.com/allenai/olmocr/blob/main/scripts/plots/ocr_pareto.png?raw=true" width=800/>
|
||||
</div>
|
||||
|
||||
## Results
|
||||
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th></th>
|
||||
<th>ArXiv</th>
|
||||
<th>Old<br>scans<br>math</th>
|
||||
<th>Tables</th>
|
||||
<th>Old<br>scans</th>
|
||||
<th>Headers<br>&<br>footers</th>
|
||||
<th>Multi<br>column</th>
|
||||
<th>Long<br>tiny<br>text</th>
|
||||
<th>Base</th>
|
||||
<th>Overall</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>Mistral OCR API</td>
|
||||
<td>77.2</td>
|
||||
<td>67.5</td>
|
||||
<td>60.6</td>
|
||||
<td>29.3</td>
|
||||
<td>93.6</td>
|
||||
<td>71.3</td>
|
||||
<td>77.1</td>
|
||||
<td>99.4</td>
|
||||
<td>72.0±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Marker 1.10.1</td>
|
||||
<td>83.8</td>
|
||||
<td>66.8</td>
|
||||
<td>72.9</td>
|
||||
<td>33.5</td>
|
||||
<td>86.6</td>
|
||||
<td>80.0</td>
|
||||
<td>85.7</td>
|
||||
<td>99.3</td>
|
||||
<td>76.1±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>MinerU 2.5.4*</td>
|
||||
<td>76.6</td>
|
||||
<td>54.6</td>
|
||||
<td>84.9</td>
|
||||
<td>33.7</td>
|
||||
<td>96.6</td>
|
||||
<td>78.2</td>
|
||||
<td>83.5</td>
|
||||
<td>93.7</td>
|
||||
<td>75.2±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>DeepSeek-OCR</td>
|
||||
<td>77.2</td>
|
||||
<td>73.6</td>
|
||||
<td>80.2</td>
|
||||
<td>33.3</td>
|
||||
<td>96.1</td>
|
||||
<td>66.4</td>
|
||||
<td>79.4</td>
|
||||
<td>99.8</td>
|
||||
<td>75.7±1.0</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Nanonets-OCR2-3B</td>
|
||||
<td>75.4</td>
|
||||
<td>46.1</td>
|
||||
<td>86.8</td>
|
||||
<td>40.9</td>
|
||||
<td>32.1</td>
|
||||
<td>81.9</td>
|
||||
<td>93.0</td>
|
||||
<td>99.6</td>
|
||||
<td>69.5±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>PaddleOCR-VL*</td>
|
||||
<td>85.7</td>
|
||||
<td>71.0</td>
|
||||
<td>84.1</td>
|
||||
<td>37.8</td>
|
||||
<td>97.0</td>
|
||||
<td>79.9</td>
|
||||
<td>85.7</td>
|
||||
<td>98.5</td>
|
||||
<td>80.0±1.0</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Infinity-Parser 7B*</td>
|
||||
<td>84.4</td>
|
||||
<td>83.8</td>
|
||||
<td>85.0</td>
|
||||
<td>47.9</td>
|
||||
<td>88.7</td>
|
||||
<td>84.2</td>
|
||||
<td>86.4</td>
|
||||
<td>99.8</td>
|
||||
<td>82.5±?</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>Chandra OCR 0.1.0*</td>
|
||||
<td>82.2</td>
|
||||
<td>80.3</td>
|
||||
<td>88.0</td>
|
||||
<td>50.4</td>
|
||||
<td>90.8</td>
|
||||
<td>81.2</td>
|
||||
<td>92.3</td>
|
||||
<td>99.9</td>
|
||||
<td>83.1±0.9</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td colspan="10"><hr></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>olmOCR (first release)</td>
|
||||
<td>63.3</td>
|
||||
<td>67.5</td>
|
||||
<td>62.3</td>
|
||||
<td>38.6</td>
|
||||
<td>93.4</td>
|
||||
<td>67.6</td>
|
||||
<td>54.8</td>
|
||||
<td>97.9</td>
|
||||
<td>68.2±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>v0.1.60 + Dynamic temp scaling</td>
|
||||
<td>71.4</td>
|
||||
<td>73.1</td>
|
||||
<td>65.6</td>
|
||||
<td>40.5</td>
|
||||
<td>93.2</td>
|
||||
<td>76.6</td>
|
||||
<td>64.9</td>
|
||||
<td>96.7</td>
|
||||
<td>72.8±1.2</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>v0.1.68 + Better prompting</td>
|
||||
<td>76.3</td>
|
||||
<td>76.0</td>
|
||||
<td>70.2</td>
|
||||
<td>43.2</td>
|
||||
<td>94.1</td>
|
||||
<td>77.5</td>
|
||||
<td>71.9</td>
|
||||
<td>96.8</td>
|
||||
<td>75.8±1.0</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>v0.2.0 + New trainer, YAML, img resize, Qwen 2.5 VL</td>
|
||||
<td>78.8</td>
|
||||
<td>77.5</td>
|
||||
<td>71.9</td>
|
||||
<td>45.4</td>
|
||||
<td>94.2</td>
|
||||
<td>78.6</td>
|
||||
<td>81.4</td>
|
||||
<td>99.8</td>
|
||||
<td>78.5±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>v0.3.0 + Handle blank pages</td>
|
||||
<td>78.6</td>
|
||||
<td>79.9</td>
|
||||
<td>72.9</td>
|
||||
<td>43.9</td>
|
||||
<td>95.1</td>
|
||||
<td>77.3</td>
|
||||
<td>81.2</td>
|
||||
<td>98.9</td>
|
||||
<td>78.5±1.1</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>v0.4.0 + Synth data, RLVR, souping</td>
|
||||
<td>83.0</td>
|
||||
<td>82.3</td>
|
||||
<td>84.9</td>
|
||||
<td>47.7</td>
|
||||
<td>96.1</td>
|
||||
<td>83.7</td>
|
||||
<td>81.9</td>
|
||||
<td>99.7</td>
|
||||
<td>82.4±1.1</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
<sup><sub>Results are reproduced in-house, except those marked with *, which are reported by model authors.
|
||||
</sub></sup>
|
||||
|
||||
## Sourcing Documents and Tests
|
||||
|
||||
We define 7 distinct document types that we found olmOCR (or its earlier iterations) often struggled to process and defined custom acquisition strategies for each (described below). We removed documents that both contained PII and were not meant for public dissemination. We also decontaminate against documents that appear in olmOCR-Mix via URL level deduplication. To scale creation of test cases over these documents, we combined manual design and review with prompting GPT-4o.
|
||||
|
||||
### Document Types
|
||||
|
||||
- **arXiv Math (AR)**: We downloaded a recent set of papers from the math subset of arXiv, selecting manuscripts with a single TeX source file and corresponding rendered PDF. To select a candidate LATEX expression from a page to use in a test, we (1) ran olmOCR to identify candidate pages with TeX, (2) match pages back to original TeX source, and (3) validate matched TeX rendering compatibility with KaTeX. We manually verify the final set of test cases to exclude instances where custom macros produce renderings that deviate from standard LATEX and to split multi-part equations into smaller test cases.
|
||||
|
||||
- **Old Scans Math (OSM)**: We crawl old, public domain math textbooks from the Internet Archive, extracting random pages from these documents. We similarly use olmOCR to find candidate pages with formulas, but this time manually annotate each formula on the page to use as test cases.
|
||||
|
||||
- **Tables (TA)**: We sampled more documents from the same internal crawled PDF repository used to create olmOCR-Mix and filtered to those which had tables using a simple prompt with Gemini-Flash-2.0. On pages with tables, we prompted Gemini-Flash-2.0 for the relationships between randomly chosen cells. We manually reviewed those tests for accuracy.
|
||||
|
||||
- **Old Scans (OS)**: We sampled historical letters and typewritten documents with existing human transcriptions from the Library of Congress digital archives. We then wrote a small script to generate Natural Reading Order cases consisting of sentences that were naturally before or after one another in the original human transcriptions. We manually added test cases to cover some headers/footers which should have been excluded from any OCR version of these documents. All of the test cases then underwent a second pass of human review for accuracy.
|
||||
|
||||
- **Headers Footers (HF)**: We sampled documents from the same internally crawled PDF repository as olmOCR-Mix. We used DocLayout-YOLO to identify page regions labeled as headers or footers using the abandon category. To extract the text from these header/footer regions, we visually mask out the rest of the document and prompt Gemini-Flash-2.0 for the content. These extracted snippets are added as test cases that should be absent in linearized output. We manually reviewed to remove mistakenly filtered text and to set conditions such as limiting the search area to the first N or last N characters.
|
||||
|
||||
- **Multi Column (MC)**: We visually sample documents from our internal crawled PDF repository to find documents with multi-column layouts and multiple articles on one page. We use Claude-Sonnet-3.7 to render those pages to HTML, and from that HTML, we extract text segments before/after one another. We manually review each entry for accuracy. We purposely select simple text blocks from coherent regions of the document, and avoid including any math formulas, superscripts, or subscripts in these tests.
|
||||
|
||||
- **Long Tiny Text (LTT)**: We crawled documents from the Internet Archive containing a large amount of dense, small print on a single page. Such documents include pages from a dictionary or pages of references from academic papers. We then generate test cases using Gemini-Flash-2.0 and verify them manually.
|
||||
|
||||
## Benchmark Principles
|
||||
|
||||
As we created olmOCR-bench, we also kept a few general rules in mind:
|
||||
|
||||
- We expect your OCR system to output a plain-text Unicode document in a reading order that would be considered natural.
|
||||
- Documents from the benchmark should fit on a standard A4 piece of paper and still be readable to a human.
|
||||
- Markdown syntax is allowed, but ignored. Ex. if we are looking for the word "enlightenment" to appear on a page, and your system outputs "**\*\*enlightenment\*\***" in Markdown bold, that still counts.
|
||||
- olmOCR-bench is not position sensitive, ex. we check that a sentence or math equation appears anywhere on a page. The exception to this is header/footer tests where we want to find simple page numbers appearing in the first or last few characters of a page.
|
||||
- Tables can be in either Markdown syntax, or as an html `<table>`.
|
||||
- Math equations must render with [Katex](https://katex.org/) and be delimeted with $, $$, \\(, or \\[.
|
||||
- Math equations are not position sensitive either, so if we are checking for
|
||||
$ 3x^2 $ to appear on a page, then outputting $ \int_a^b{ 3x ^ 2dx} $ counts.
|
||||
- We normalize all Unicode to NFC before running the benchmark, so if your OCR model outputs é vs e + ◌́ then either way should not affect your benchmark score.
|
||||
- We normalize all the different variants of hyphens to the ascii -, all the variants of double quoets to ascii " and all variants of single quotes/apostrophes to ascii '. You should score the same on the benchmark if you output - vs —
|
||||
- All facts checked about documents are either pass/fail. We want it to be very clear if your OCR system fails a test, and if so, what output would make it pass.
|
||||
|
||||
|
||||
## olmOCR-Bench Test classes
|
||||
|
||||
- Text presence
|
||||
- This task makes sure that a given small piece of text (ex. 1-3 sentence level) is present within
|
||||
a parsed document. Soft/fuzzy matching is allowed, as well as specifying if the text must be in the first N or last N characters of the document. Case sensitive by default.
|
||||
- Text absense
|
||||
- This task makes sure that a given piece of next does NOT appear in the OCR'ed version of a document. We generally want our OCR systems to filter out content like headers/footers/page numbers from documents. The same fuzzy matching as in Text Presence tests is allowed.
|
||||
- Natural Reading Order
|
||||
- This task ensures that blocks of text which are present have a defined order relative to one another. For example,
|
||||
on a document that contains multiple news articles on one page, you'd want to see that the first sentence of the
|
||||
first article appears after the heading of that article. But, you may be okay with swapping the order of those
|
||||
two articles.
|
||||
- Table Accuracy
|
||||
- Both Markdown and HTML based tables are supported. These tests check that a cell with a given text exists somewhere in the table, and that its neighbors have certain properties. Ex. A cell exists on this page with text "4.5%" and above that is a cell with the text "2.4%". However, it's important to note that some tests depend on rowspan and colspan information being present in the table, which is only available with HTML based tables. This means that a model outputting only markdown tables cannot achieve a max score on this section.
|
||||
- Math Formula Accuracy
|
||||
- We render a given Latex style equation using Katex in a headless browser. And then see if it exists anywhere in the final OCRed document. Matching is performed on a relative symbol level, ex. in "\f\relax{x} = \int_{-\infty}^\infty
|
||||
x^2dx" we check that a ∫ appears to the left of a x, x appears to the left of dx, etc...
|
||||
|
||||
|
||||
|
||||
## Downloading and running the benchmark
|
||||
|
||||
Currently the full benchmark data is located here:
|
||||
https://huggingface.co/datasets/allenai/olmOCR-bench
|
||||
|
||||
To run a benchmark, first install the bench requirements
|
||||
```bash
|
||||
conda create -n olmocr python=3.11
|
||||
conda activate olmocr
|
||||
|
||||
git clone https://github.com/allenai/olmocr.git
|
||||
cd olmocr
|
||||
|
||||
# Install olmocr and the requirements needed to run the benchmark
|
||||
pip install -e .[bench]
|
||||
|
||||
# Configure playwright headless browser to run the math rendering tests
|
||||
playwright install chromium
|
||||
|
||||
# Now clone the benchmark data from hugging face, this includes the PDFs and JSON annotation data
|
||||
huggingface-cli download --repo-type dataset --resume-download allenai/olmOCR-bench --local-dir ./olmOCR-bench
|
||||
```
|
||||
|
||||
Convert your documents
|
||||
```bash
|
||||
# You will need to install the [gpu] subset of olmocr dependencies to run gpu inference
|
||||
# Then convert using using olmocr.bench.convert, see the olmocr/bench/runners directory for options
|
||||
pip install olmocr[gpu] --find-links https://flashinfer.ai/whl/cu124/torch2.4/flashinfer/
|
||||
python -m olmocr.bench.convert olmocr_pipeline --dir ./olmOCR-bench/bench_data
|
||||
|
||||
# OR, you can use the pipeline to convert the benchmark PDFs and move them into the final format
|
||||
python -m olmocr.pipeline ./localworkspace --markdown --pdfs ./olmOCR-bench/bench_data/pdfs/**/*.pdf
|
||||
python olmocr/bench/scripts/workspace_to_bench.py localworkspace/ olmOCR-bench/bench_data/olmocr --bench-path ./olmOCR-bench/
|
||||
```
|
||||
|
||||
Now run the benchmark
|
||||
```bash
|
||||
python -m olmocr.bench.benchmark --dir ./olmOCR-bench/bench_data
|
||||
```
|
||||
|
||||
## Previewing the benchmark questions
|
||||
|
||||
We have an internal data annotation tool that can be used to review the questions in the benchmark, and make edits.
|
||||
|
||||
<img width="700" alt="image" src="https://github.com/user-attachments/assets/dd24fd88-a642-4379-b5a1-9911717bf5b1" />
|
||||
|
||||
|
||||
```bash
|
||||
python -m olmocr.bench.review_app --port 5000 --debug ./olmOCR-bench/bench_data/multi_column.jsonl --force
|
||||
```
|
||||
@@ -0,0 +1,463 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
This script runs olmocr bench.
|
||||
It will take as an argument a folder, and scan it for .jsonl files which contain the various rules and properties that we will check.
|
||||
It will then validate the JSON files to make sure they are all valid.
|
||||
Then, each other folder in there (besides /pdfs) represents a pipeline tool that we will evaluate.
|
||||
We will validate that each one of those contains at least one .md file (or repeated generations, e.g. _pg{page}_repeat{repeat}.md)
|
||||
corresponding to its parse for every .pdf in the /pdfs folder.
|
||||
Then, we will read each one, and check if they pass against all the rules.
|
||||
If a rule fails on some of the repeats, a short explanation is printed.
|
||||
The final score is the average of per-JSONL file scores, where each JSONL file's score is the proportion of tests from that file that pass.
|
||||
Statistical analysis including bootstrap confidence intervals are provided for the results.
|
||||
Pairwise permutation tests are conducted between specific candidate pairs.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import sys
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
from pypdf import PdfReader
|
||||
from tqdm import tqdm
|
||||
|
||||
from .report import generate_html_report
|
||||
from .tests import BaselineTest, BasePDFTest, load_tests, save_tests
|
||||
from .utils import calculate_bootstrap_ci
|
||||
|
||||
|
||||
def evaluate_candidate(
|
||||
candidate_folder: str, all_tests: List[BasePDFTest], pdf_basenames: List[str], force: bool = False
|
||||
) -> Tuple[float, int, List[str], List[str], Dict[str, List[float]], List[float], Dict[str, Dict[int, List[Tuple[BasePDFTest, bool, str]]]]]:
|
||||
"""
|
||||
For the candidate folder (pipeline tool output), validate that it contains at least one .md file
|
||||
(i.e. repeated generations like _pg{page}_repeat{repeat}.md) for every PDF in the pdf folder.
|
||||
Then, run each rule against all corresponding .md files concurrently and average the results.
|
||||
|
||||
Returns a tuple:
|
||||
(overall_score, total_tests, candidate_errors, test_failures, test_type_breakdown, all_test_scores, test_results)
|
||||
|
||||
- overall_score: Average fraction of tests passed (averaged over repeats and tests).
|
||||
Note: This is now updated at reporting time to be the average of per-JSONL file scores.
|
||||
- total_tests: Total number of tests evaluated.
|
||||
- candidate_errors: List of candidate errors (e.g. missing files).
|
||||
- test_failures: List of failure messages for tests not passing on all repeats.
|
||||
- test_type_breakdown: Dictionary mapping test type to list of average pass ratios for tests of that type.
|
||||
- all_test_scores: List of all individual test scores (used for bootstrapping).
|
||||
- test_results: Dictionary mapping PDF name to dictionary mapping page number to list of (test, passed, explanation) tuples.
|
||||
"""
|
||||
candidate_errors = []
|
||||
test_failures = []
|
||||
test_type_breakdown = {} # key: test type, value: list of average pass ratios
|
||||
all_test_scores = [] # Store all individual test scores for bootstrapping
|
||||
test_results = {} # Store detailed test results for reporting
|
||||
candidate_name = os.path.basename(candidate_folder)
|
||||
|
||||
# Map each PDF to its corresponding MD repeats (e.g., doc1_pg1_repeat1.md, doc1_pg2_repeat2.md, etc.)
|
||||
pdf_to_md_files = {}
|
||||
all_files = list(glob.glob(os.path.join(candidate_folder, "**/*.md"), recursive=True))
|
||||
|
||||
for pdf_name in pdf_basenames:
|
||||
md_base = os.path.splitext(pdf_name)[0]
|
||||
md_regex = re.compile(rf"^{re.escape(md_base)}_pg\d+_repeat\d+\.md$")
|
||||
md_files = [f for f in all_files if md_regex.match(os.path.relpath(f, candidate_folder))]
|
||||
|
||||
if not md_files and not force:
|
||||
candidate_errors.append(
|
||||
f"Candidate '{candidate_name}' is missing MD repeats for {pdf_name} " f"(expected files matching {md_base}_pg{{page}}_repeat*.md)."
|
||||
)
|
||||
else:
|
||||
pdf_to_md_files[pdf_name] = md_files
|
||||
|
||||
if candidate_errors:
|
||||
return (0.0, len(all_tests), candidate_errors, test_failures, test_type_breakdown, all_test_scores, test_results)
|
||||
|
||||
# Define an inner function to evaluate a single test
|
||||
def process_test(test: BasePDFTest) -> Tuple[float, str, str, List[str], Tuple[bool, str]]:
|
||||
local_errors = []
|
||||
test_failure = None
|
||||
pdf_name = test.pdf
|
||||
|
||||
# Initialize the test_results structure if needed
|
||||
if pdf_name not in test_results:
|
||||
test_results[pdf_name] = {}
|
||||
if test.page not in test_results[pdf_name]:
|
||||
test_results[pdf_name][test.page] = []
|
||||
|
||||
md_base = os.path.splitext(pdf_name)[0]
|
||||
md_files = pdf_to_md_files.get(pdf_name, [])
|
||||
# Filter MD files for the specific page corresponding to the test
|
||||
page_md_files = [f for f in md_files if re.search(rf"_pg{test.page}_", os.path.basename(f))]
|
||||
if not page_md_files:
|
||||
local_errors.append(
|
||||
f"Candidate '{candidate_name}' is missing MD repeats for {pdf_name} page {test.page} "
|
||||
f"(expected files matching {md_base}_pg{test.page}_repeat*.md)."
|
||||
)
|
||||
test_results[pdf_name][test.page].append((test, False, "Missing MD files"))
|
||||
return (0.0, None, test.type, local_errors, (False, "Missing MD files"))
|
||||
|
||||
repeat_passes = 0
|
||||
num_repeats = 0
|
||||
explanations = []
|
||||
for md_path in page_md_files:
|
||||
num_repeats += 1
|
||||
try:
|
||||
with open(md_path, "r", encoding="utf-8") as f:
|
||||
md_content = f.read()
|
||||
except Exception as e:
|
||||
local_errors.append(f"Error reading {md_path}: {e}")
|
||||
continue
|
||||
|
||||
try:
|
||||
passed, explanation = test.run(md_content)
|
||||
if passed:
|
||||
repeat_passes += 1
|
||||
else:
|
||||
explanations.append(explanation)
|
||||
except Exception as e:
|
||||
local_errors.append(f"Error running test {test.id} on {md_path}: {e}")
|
||||
explanations.append(str(e))
|
||||
|
||||
test_avg = repeat_passes / num_repeats if num_repeats > 0 else 0.0
|
||||
final_passed = test_avg > 0.5 # Consider test passed if majority of repeats pass
|
||||
final_explanation = explanations[0] if explanations else "All repeats passed"
|
||||
|
||||
# Store the test result for reporting
|
||||
test_results[pdf_name][test.page].append((test, final_passed, final_explanation))
|
||||
|
||||
if test_avg < 1.0:
|
||||
test_failure = (
|
||||
f"Test {test.id} on {md_base} page {test.page} average pass ratio: {test_avg:.3f} "
|
||||
f"({repeat_passes}/{num_repeats} repeats passed). Ex: {explanations[0] if explanations else 'No explanation'}"
|
||||
)
|
||||
return (test_avg, test_failure, test.type, local_errors, (final_passed, final_explanation))
|
||||
|
||||
total_test_score = 0.0
|
||||
futures = []
|
||||
# Use a thread pool to evaluate each test concurrently.
|
||||
with ThreadPoolExecutor(max_workers=min(os.cpu_count() or 1, 64)) as executor:
|
||||
futures = [executor.submit(process_test, test) for test in all_tests]
|
||||
# tqdm progress bar for this candidate's tests
|
||||
for future in tqdm(as_completed(futures), total=len(futures), desc=f"Evaluating tests for {candidate_name}", unit="test"):
|
||||
test_avg, test_failure, test_type, errors, _ = future.result()
|
||||
all_test_scores.append(test_avg)
|
||||
total_test_score += test_avg
|
||||
if test_failure:
|
||||
test_failures.append(test_failure)
|
||||
if test_type not in test_type_breakdown:
|
||||
test_type_breakdown[test_type] = []
|
||||
test_type_breakdown[test_type].append(test_avg)
|
||||
local_errors = errors
|
||||
if local_errors:
|
||||
candidate_errors.extend(local_errors)
|
||||
|
||||
overall_score = total_test_score / len(all_tests) if all_tests else 0.0
|
||||
return (overall_score, len(all_tests), candidate_errors, test_failures, test_type_breakdown, all_test_scores, test_results)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Run OLMOCR Bench.")
|
||||
parser.add_argument(
|
||||
"--dir",
|
||||
default=os.path.join(os.path.dirname(__file__), "sample_data"),
|
||||
help="Path to the folder containing .jsonl files, /pdfs folder, and pipeline tool subfolders.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help="Run benchmark even if some files are missing",
|
||||
)
|
||||
parser.add_argument("--candidate", type=str, default=None, help="Run test only for a single candidate")
|
||||
parser.add_argument("--skip_baseline", action="store_true", help="Skip running baseline tests (ex. that check that basic content is present on each page)")
|
||||
parser.add_argument(
|
||||
"--bootstrap_samples",
|
||||
type=int,
|
||||
default=1000,
|
||||
help="Number of bootstrap samples for confidence interval calculation (default: 1000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_level",
|
||||
type=float,
|
||||
default=0.95,
|
||||
help="Confidence level for interval calculation (default: 0.95 for 95% CI).",
|
||||
)
|
||||
# New arguments
|
||||
parser.add_argument("--sample", type=int, default=None, help="Randomly sample N tests to run instead of all tests.")
|
||||
parser.add_argument("--test_report", type=str, default=None, help="Generate an HTML report of test results. Provide a filename (e.g., results.html).")
|
||||
parser.add_argument(
|
||||
"--output_failed", type=str, default=None, help="Output a JSONL file containing tests that failed across all candidates. Provide a filename."
|
||||
)
|
||||
parser.add_argument("--max_reports", type=int, default=None, help="Limit the HTML report to at most N unique PDFs per .jsonl file.")
|
||||
args = parser.parse_args()
|
||||
|
||||
input_folder = args.dir if os.path.isdir(args.dir) else os.path.dirname(args.dir)
|
||||
n_bootstrap = args.bootstrap_samples
|
||||
ci_level = args.confidence_level
|
||||
pdf_folder = os.path.join(input_folder, "pdfs")
|
||||
|
||||
if not os.path.exists(pdf_folder):
|
||||
print("Error: /pdfs folder must exist in your data directory.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
all_pdf_files = list(glob.glob(os.path.join(pdf_folder, "**/*.pdf"), recursive=True))
|
||||
|
||||
if not all_pdf_files:
|
||||
print(f"Error: No PDF files found in {pdf_folder}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
pdf_basenames = [os.path.relpath(p, pdf_folder) for p in all_pdf_files]
|
||||
|
||||
if os.path.isfile(args.dir):
|
||||
jsonl_files = [args.dir]
|
||||
else:
|
||||
jsonl_files = glob.glob(os.path.join(input_folder, "*.jsonl"))
|
||||
|
||||
if not jsonl_files:
|
||||
print(f"Error: No .jsonl files found in {input_folder}.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
all_tests = []
|
||||
test_to_jsonl = {} # Map test IDs to their source jsonl files
|
||||
for jsonl_path in jsonl_files:
|
||||
jsonl_basename = os.path.basename(jsonl_path)
|
||||
tests = load_tests(jsonl_path)
|
||||
for test in tests:
|
||||
test_to_jsonl[test.id] = jsonl_basename
|
||||
all_tests.extend(tests)
|
||||
|
||||
if not all_tests:
|
||||
print("No valid tests found. Exiting.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# When a single .jsonl file is passed, only consider PDFs referenced by its tests
|
||||
if os.path.isfile(args.dir):
|
||||
referenced_pdfs = {t.pdf for t in all_tests}
|
||||
pdf_basenames = [p for p in pdf_basenames if p in referenced_pdfs]
|
||||
|
||||
for pdf in pdf_basenames:
|
||||
if not any(t.type == "baseline" for t in all_tests if t.pdf == pdf):
|
||||
all_tests.append(BaselineTest(id=f"{pdf}_baseline", pdf=pdf, page=1, type="baseline"))
|
||||
test_to_jsonl[all_tests[-1].id] = "baseline"
|
||||
|
||||
for pdf in pdf_basenames:
|
||||
pdf_doc = PdfReader(os.path.join(pdf_folder, pdf))
|
||||
for page in range(1, len(pdf_doc.pages) + 1):
|
||||
if not any(test for test in all_tests if test.pdf == pdf and test.page == page) and not args.force:
|
||||
print(f"No dataset entry found for pdf {pdf} page {page}")
|
||||
sys.exit(1)
|
||||
|
||||
if args.skip_baseline:
|
||||
all_tests = [test for test in all_tests if test.type != "baseline"]
|
||||
|
||||
# Sample tests if requested
|
||||
if args.sample is not None and args.sample > 0:
|
||||
if args.sample >= len(all_tests):
|
||||
print(f"Sample size {args.sample} is greater than or equal to the total number of tests ({len(all_tests)}). Using all tests.")
|
||||
else:
|
||||
print(f"Randomly sampling {args.sample} tests out of {len(all_tests)} total tests.")
|
||||
all_tests = random.sample(all_tests, args.sample)
|
||||
|
||||
candidate_folders = []
|
||||
for entry in os.listdir(input_folder):
|
||||
full_path = os.path.join(input_folder, entry)
|
||||
if args.candidate is not None:
|
||||
if entry == args.candidate:
|
||||
candidate_folders.append(full_path)
|
||||
else:
|
||||
if os.path.isdir(full_path) and entry != "pdfs":
|
||||
candidate_folders.append(full_path)
|
||||
|
||||
if not candidate_folders:
|
||||
print("Error: No candidate pipeline folders found (subdirectories besides 'pdfs').", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
candidate_folders.sort()
|
||||
|
||||
summary = []
|
||||
test_results_by_candidate = {}
|
||||
print("\nRunning tests for each candidate:")
|
||||
# Process candidates sequentially so that each candidate's progress bar is distinct.
|
||||
for candidate in candidate_folders:
|
||||
candidate_name = os.path.basename(candidate)
|
||||
print(f"\nEvaluating candidate: {candidate_name}")
|
||||
overall_score, total_tests, candidate_errors, test_failures, test_type_breakdown, all_test_scores, test_results = evaluate_candidate(
|
||||
candidate, all_tests, pdf_basenames, args.force
|
||||
)
|
||||
|
||||
# Always store test results for displaying jsonl file groupings
|
||||
test_results_by_candidate[candidate_name] = test_results
|
||||
|
||||
# Group results by jsonl file for more accurate CI calculation
|
||||
jsonl_results = {}
|
||||
jsonl_scores = [] # List to store scores by jsonl file for CI calculation
|
||||
jsonl_file_sizes = [] # List to store the number of tests per jsonl file
|
||||
|
||||
for test in all_tests:
|
||||
# Get the jsonl file this test came from
|
||||
jsonl_file = test_to_jsonl.get(test.id, "unknown")
|
||||
|
||||
if jsonl_file not in jsonl_results:
|
||||
jsonl_results[jsonl_file] = {"total": 0, "passed": 0, "scores": []}
|
||||
|
||||
jsonl_results[jsonl_file]["total"] += 1
|
||||
|
||||
# Get the test result for this candidate if it exists
|
||||
if not candidate_errors and hasattr(test, "pdf") and hasattr(test, "page"):
|
||||
pdf_name = test.pdf
|
||||
page = test.page
|
||||
if pdf_name in test_results and page in test_results.get(pdf_name, {}):
|
||||
for t, passed, _ in test_results[pdf_name][page]:
|
||||
if t.id == test.id:
|
||||
# Store the test score in its jsonl group
|
||||
result_score = 1.0 if passed else 0.0
|
||||
jsonl_results[jsonl_file]["scores"].append(result_score)
|
||||
if passed:
|
||||
jsonl_results[jsonl_file]["passed"] += 1
|
||||
break
|
||||
|
||||
# Gather all the scores by jsonl file for CI calculation
|
||||
for jsonl_file, results in jsonl_results.items():
|
||||
if results["scores"]:
|
||||
jsonl_file_sizes.append(len(results["scores"]))
|
||||
jsonl_scores.extend(results["scores"])
|
||||
|
||||
# Calculate CI using the updated function with splits
|
||||
if jsonl_scores:
|
||||
ci = calculate_bootstrap_ci(jsonl_scores, n_bootstrap=n_bootstrap, ci_level=ci_level, splits=jsonl_file_sizes)
|
||||
else:
|
||||
ci = (0.0, 0.0)
|
||||
summary.append((candidate_name, overall_score, total_tests, candidate_errors, test_failures, test_type_breakdown, ci, all_test_scores))
|
||||
print(f"\nCandidate: {candidate_name}")
|
||||
if candidate_errors:
|
||||
for err in candidate_errors:
|
||||
print(f" [ERROR] {err}")
|
||||
else:
|
||||
if test_failures:
|
||||
for fail in test_failures:
|
||||
print(f" [FAIL] {fail}")
|
||||
# Calculate and show the per-category average score
|
||||
jsonl_pass_rates = []
|
||||
for _, results in jsonl_results.items():
|
||||
if results["total"] > 0:
|
||||
pass_rate = results["passed"] / results["total"]
|
||||
jsonl_pass_rates.append(pass_rate)
|
||||
|
||||
per_category_score = sum(jsonl_pass_rates) / len(jsonl_pass_rates) if jsonl_pass_rates else 0.0
|
||||
print(f" Average Score: {per_category_score * 100:.1f}% (95% CI: [{ci[0] * 100:.1f}%, {ci[1] * 100:.1f}%]) over {total_tests} tests.")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Final Summary with 95% Confidence Intervals:")
|
||||
for idx, (candidate_name, _, total_tests, candidate_errors, _, test_type_breakdown, ci, _) in enumerate(summary):
|
||||
# Group results by jsonl file
|
||||
jsonl_results = {}
|
||||
for test in all_tests:
|
||||
# Get the jsonl file this test came from
|
||||
jsonl_file = test_to_jsonl.get(test.id, "unknown")
|
||||
|
||||
if jsonl_file not in jsonl_results:
|
||||
jsonl_results[jsonl_file] = {"total": 0, "passed": 0}
|
||||
|
||||
jsonl_results[jsonl_file]["total"] += 1
|
||||
|
||||
# Get the test result for this candidate if it exists
|
||||
test_result = None
|
||||
if not candidate_errors and hasattr(test, "pdf") and hasattr(test, "page"):
|
||||
pdf_name = test.pdf
|
||||
page = test.page
|
||||
if pdf_name in test_results_by_candidate.get(candidate_name, {}) and page in test_results_by_candidate[candidate_name].get(pdf_name, {}):
|
||||
for t, passed, _ in test_results_by_candidate[candidate_name][pdf_name][page]:
|
||||
if t.id == test.id:
|
||||
test_result = passed
|
||||
break
|
||||
|
||||
if test_result:
|
||||
jsonl_results[jsonl_file]["passed"] += 1
|
||||
|
||||
# Calculate new overall score as average of per-JSONL pass rates
|
||||
jsonl_pass_rates = []
|
||||
for jsonl_file, results in jsonl_results.items():
|
||||
if results["total"] > 0:
|
||||
pass_rate = results["passed"] / results["total"]
|
||||
jsonl_pass_rates.append(pass_rate)
|
||||
|
||||
# New overall score is average of per-JSONL pass rates
|
||||
new_overall_score = sum(jsonl_pass_rates) / len(jsonl_pass_rates) if jsonl_pass_rates else 0.0
|
||||
|
||||
# Update the overall_score in the summary list for later use (e.g., in permutation tests)
|
||||
summary[idx] = (candidate_name, new_overall_score, total_tests, candidate_errors, summary[idx][4], test_type_breakdown, ci, summary[idx][7])
|
||||
|
||||
if candidate_errors:
|
||||
status = "FAILED (errors)"
|
||||
ciw_str = ""
|
||||
else:
|
||||
status = f"{new_overall_score * 100:0.1f}%"
|
||||
# Use the CI that was calculated with proper category-based bootstrap
|
||||
half_width = ((ci[1] - ci[0]) / 2) * 100
|
||||
ciw_str = f"± {half_width:0.1f}%"
|
||||
print(f"{candidate_name:20s} : Average Score: {status} {ciw_str} (average of per-JSONL scores)")
|
||||
|
||||
# Sort the test types alphabetically
|
||||
for ttype in sorted(test_type_breakdown.keys()):
|
||||
scores = test_type_breakdown[ttype]
|
||||
avg = sum(scores) / len(scores) * 100 if scores else 0.0
|
||||
print(f" {ttype:8s}: {avg:0.1f}% average pass rate over {len(scores)} tests")
|
||||
|
||||
print("\n Results by JSONL file:")
|
||||
for jsonl_file, results in sorted(jsonl_results.items()):
|
||||
if results["total"] > 0:
|
||||
pass_rate = (results["passed"] / results["total"]) * 100
|
||||
print(f" {jsonl_file:30s}: {pass_rate:0.1f}% ({results['passed']}/{results['total']} tests)")
|
||||
print("")
|
||||
|
||||
# Generate HTML report if requested
|
||||
if args.test_report:
|
||||
generate_html_report(test_results_by_candidate, pdf_folder, args.test_report, max_reports=args.max_reports, test_to_jsonl=test_to_jsonl)
|
||||
|
||||
# Output tests that failed across all candidates if requested
|
||||
if args.output_failed:
|
||||
# Identify tests that failed across all candidates
|
||||
all_failed_tests = []
|
||||
valid_candidates = [c for c in summary if not c[3]] # Skip candidates with errors
|
||||
|
||||
for test in all_tests:
|
||||
# Track whether this test has any results
|
||||
has_results = False
|
||||
any_passed = False
|
||||
|
||||
for candidate_name, _, _, _, _, _, _, _ in valid_candidates:
|
||||
# Get the test result for this candidate
|
||||
test_result = None
|
||||
if hasattr(test, "pdf") and hasattr(test, "page"):
|
||||
pdf_name = test.pdf
|
||||
page = test.page
|
||||
if pdf_name in test_results_by_candidate.get(candidate_name, {}) and page in test_results_by_candidate[candidate_name].get(pdf_name, {}):
|
||||
for t, passed, explanation in test_results_by_candidate[candidate_name][pdf_name][page]:
|
||||
if t.id == test.id:
|
||||
has_results = True
|
||||
test_result = passed
|
||||
if passed:
|
||||
any_passed = True
|
||||
break
|
||||
|
||||
# If we have results for this test and it never passed for any candidate, add it to the failed list
|
||||
if has_results and not any_passed:
|
||||
# Add to the list
|
||||
all_failed_tests.append(test)
|
||||
|
||||
# If we have any failed tests, write them to the specified JSONL file
|
||||
output_path = os.path.join(input_folder, args.output_failed) if not os.path.isabs(args.output_failed) else args.output_failed
|
||||
|
||||
if all_failed_tests:
|
||||
save_tests(all_failed_tests, output_path)
|
||||
|
||||
print(f"\nOutput {len(all_failed_tests)} tests that failed across all candidates to {output_path}")
|
||||
else:
|
||||
print("\nNo tests failed across all candidates. No output file created.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,305 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import glob
|
||||
import importlib
|
||||
import os
|
||||
import tempfile
|
||||
import traceback
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from functools import partial
|
||||
|
||||
from pypdf import PdfReader
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.image_utils import convert_image_to_pdf_bytes
|
||||
|
||||
|
||||
def parse_method_arg(method_arg):
|
||||
"""
|
||||
Parse a method configuration string of the form:
|
||||
method_name[:key=value[:key2=value2...]]
|
||||
Returns:
|
||||
(method_name, kwargs_dict, folder_name)
|
||||
"""
|
||||
parts = method_arg.split(":")
|
||||
name = parts[0]
|
||||
kwargs = {}
|
||||
folder_name = name # Default folder name is the method name
|
||||
|
||||
for extra in parts[1:]:
|
||||
if "=" in extra:
|
||||
key, value = extra.split("=", 1)
|
||||
if key == "name":
|
||||
folder_name = value
|
||||
continue
|
||||
|
||||
try:
|
||||
converted = int(value)
|
||||
except ValueError:
|
||||
try:
|
||||
converted = float(value)
|
||||
except ValueError:
|
||||
converted = value
|
||||
kwargs[key] = converted
|
||||
else:
|
||||
raise ValueError(f"Extra argument '{extra}' is not in key=value format")
|
||||
|
||||
return name, kwargs, folder_name
|
||||
|
||||
|
||||
# Wrapper to run synchronous functions in the event loop
|
||||
async def run_sync_in_executor(func, executor, *args, **kwargs):
|
||||
"""Run a synchronous function in the provided executor"""
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(executor, partial(func, *args, **kwargs))
|
||||
|
||||
|
||||
async def process_pdf(pdf_path, page_num, method, kwargs, output_path, is_async, executor=None, use_executor=True, failfast=False):
|
||||
"""Process a single PDF and save the result to output_path"""
|
||||
try:
|
||||
if is_async:
|
||||
# Run async function directly
|
||||
markdown = await method(pdf_path, page_num=page_num, **kwargs)
|
||||
elif use_executor:
|
||||
# Run synchronous function in the executor
|
||||
markdown = await run_sync_in_executor(method, executor, pdf_path, page_num=page_num, **kwargs)
|
||||
else:
|
||||
# Run synchronous function directly without executor (when parallel=0)
|
||||
markdown = method(pdf_path, page_num=page_num, **kwargs)
|
||||
|
||||
if markdown is None:
|
||||
print(f"Warning, did not get output for {os.path.basename(output_path)}")
|
||||
# Write blank to this file, so that it's marked as an error and not just skipped in evals
|
||||
with open(output_path, "w") as out_f:
|
||||
out_f.write("")
|
||||
return False
|
||||
|
||||
# Write the markdown to the output file
|
||||
with open(output_path, "w") as out_f:
|
||||
out_f.write(markdown)
|
||||
|
||||
return True
|
||||
except Exception as ex:
|
||||
if failfast:
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"FATAL ERROR: Exception occurred while processing {os.path.basename(output_path)}")
|
||||
print(f"PDF: {pdf_path}, Page: {page_num}")
|
||||
print(f"{'=' * 60}")
|
||||
print("\nFull stack trace:")
|
||||
traceback.print_exc()
|
||||
print(f"{'=' * 60}\n")
|
||||
# Re-raise the exception to stop processing
|
||||
raise
|
||||
else:
|
||||
print(f"Exception {str(ex)} occurred while processing {os.path.basename(output_path)}")
|
||||
# Write blank to this file, so that it's marked as an error and not just skipped in evals
|
||||
with open(output_path, "w") as out_f:
|
||||
out_f.write("")
|
||||
return False
|
||||
|
||||
|
||||
async def process_pdfs(config, pdf_directory, data_directory, repeats, remove_text, force, max_parallel=None, failfast=False):
|
||||
"""
|
||||
Process PDFs using asyncio for both sync and async methods,
|
||||
limiting the number of concurrent tasks to max_parallel.
|
||||
"""
|
||||
# When max_parallel is 0, run synchronously without any executor
|
||||
# When max_parallel is > 0, create a thread pool executor
|
||||
use_executor = max_parallel != 0
|
||||
executor = ThreadPoolExecutor(max_workers=max_parallel) if use_executor else None
|
||||
|
||||
try:
|
||||
for candidate in config.keys():
|
||||
print(f"Starting conversion using {candidate} with kwargs: {config[candidate]['kwargs']}")
|
||||
folder_name = config[candidate]["folder_name"]
|
||||
candidate_output_dir = os.path.join(data_directory, folder_name)
|
||||
os.makedirs(candidate_output_dir, exist_ok=True)
|
||||
|
||||
method = config[candidate]["method"]
|
||||
kwargs = config[candidate]["kwargs"]
|
||||
is_async = asyncio.iscoroutinefunction(method)
|
||||
|
||||
# Use recursive glob to support nested PDFs
|
||||
all_pdfs = glob.glob(os.path.join(pdf_directory, "**/*.pdf"), recursive=True)
|
||||
all_pdfs.sort()
|
||||
|
||||
# Prepare all tasks
|
||||
tasks = []
|
||||
task_descriptions = {}
|
||||
|
||||
for pdf_path in all_pdfs:
|
||||
pdf = PdfReader(pdf_path)
|
||||
num_pages = len(pdf.pages)
|
||||
base_name = os.path.basename(pdf_path).replace(".pdf", "")
|
||||
# Determine the PDF's relative folder path (e.g. "arxiv_data") relative to pdf_directory
|
||||
relative_pdf_path = os.path.relpath(pdf_path, pdf_directory)
|
||||
pdf_relative_dir = os.path.dirname(relative_pdf_path)
|
||||
|
||||
if remove_text:
|
||||
print(f"Converting {pdf_path} into images to remove text-content...")
|
||||
|
||||
# Generate image files from each page
|
||||
temp_image_files = []
|
||||
try:
|
||||
for page_num in range(1, num_pages + 1):
|
||||
# Get base64 PNG data for the current page
|
||||
base64_png = render_pdf_to_base64png(pdf_path, page_num, target_longest_image_dim=2048)
|
||||
|
||||
# Decode base64 and save to temporary file
|
||||
temp_img = tempfile.NamedTemporaryFile("wb", suffix=".png", delete=False)
|
||||
temp_img.write(base64.b64decode(base64_png))
|
||||
temp_img.close()
|
||||
temp_image_files.append(temp_img.name)
|
||||
|
||||
# Convert all images to a single PDF using our enhanced function
|
||||
pdf_bytes = convert_image_to_pdf_bytes(temp_image_files)
|
||||
|
||||
# Write the PDF bytes to a temporary file
|
||||
temp_pdf = tempfile.NamedTemporaryFile("wb", suffix=".pdf", delete=False)
|
||||
temp_pdf.write(pdf_bytes)
|
||||
temp_pdf.close()
|
||||
|
||||
# Update pdf_path to the new file
|
||||
pdf_path = temp_pdf.name
|
||||
|
||||
finally:
|
||||
# Clean up temporary image files
|
||||
for temp_file in temp_image_files:
|
||||
try:
|
||||
os.remove(temp_file)
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to remove temporary file {temp_file}: {e}")
|
||||
|
||||
for repeat in range(1, repeats + 1):
|
||||
for page_num in range(1, num_pages + 1):
|
||||
output_filename = f"{base_name}_pg{page_num}_repeat{repeat}.md"
|
||||
# Preserve the relative folder structure in the output directory
|
||||
candidate_pdf_dir = os.path.join(candidate_output_dir, pdf_relative_dir)
|
||||
os.makedirs(candidate_pdf_dir, exist_ok=True)
|
||||
output_path = os.path.join(candidate_pdf_dir, output_filename)
|
||||
|
||||
if os.path.exists(output_path) and not force:
|
||||
print(f"Skipping {base_name}_pg{page_num}_repeat{repeat} for {candidate}, file already exists")
|
||||
print("Rerun with --force flag to force regeneration")
|
||||
continue
|
||||
|
||||
task = process_pdf(pdf_path, page_num, method, kwargs, output_path, is_async, executor, use_executor, failfast)
|
||||
tasks.append(task)
|
||||
task_descriptions[id(task)] = f"{base_name}_pg{page_num}_repeat{repeat} ({candidate})"
|
||||
|
||||
# Process tasks with semaphore to limit concurrency
|
||||
# When max_parallel is 0, set semaphore to 1 to run sequentially
|
||||
semaphore = asyncio.Semaphore(max_parallel if max_parallel else 1)
|
||||
|
||||
async def process_with_semaphore(task):
|
||||
async with semaphore:
|
||||
return await task
|
||||
|
||||
# Wrap each task with the semaphore
|
||||
limited_tasks = [process_with_semaphore(task) for task in tasks]
|
||||
|
||||
# Process tasks with progress bar
|
||||
if limited_tasks:
|
||||
completed = 0
|
||||
with tqdm(total=len(limited_tasks), desc=f"Processing {candidate}") as pbar:
|
||||
# When parallel=0, tasks complete synchronously and we need to handle them differently
|
||||
if max_parallel == 0:
|
||||
# Process tasks sequentially with immediate progress updates
|
||||
for task in limited_tasks:
|
||||
try:
|
||||
result = await task
|
||||
if result:
|
||||
completed += 1
|
||||
except Exception as e:
|
||||
print(f"Task failed: {e}")
|
||||
finally:
|
||||
pbar.update(1)
|
||||
else:
|
||||
# Use as_completed for parallel processing
|
||||
for task in asyncio.as_completed(limited_tasks):
|
||||
try:
|
||||
result = await task
|
||||
if result:
|
||||
completed += 1
|
||||
except Exception as e:
|
||||
print(f"Task failed: {e}")
|
||||
finally:
|
||||
pbar.update(1)
|
||||
|
||||
print(f"Completed {completed} out of {len(limited_tasks)} tasks for {candidate}")
|
||||
finally:
|
||||
# Clean up the executor
|
||||
if executor:
|
||||
executor.shutdown(wait=False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Run PDF conversion using specified OCR methods and extra parameters.")
|
||||
parser.add_argument(
|
||||
"methods",
|
||||
nargs="+",
|
||||
help="Methods to run in the format method[:key=value ...]. "
|
||||
"Example: gotocr mineru:temperature=2 marker:u=3. "
|
||||
"Use 'name=folder_name' to specify a custom output folder name.",
|
||||
)
|
||||
parser.add_argument("--repeats", type=int, default=1, help="Number of times to repeat the conversion for each PDF.")
|
||||
parser.add_argument(
|
||||
"--dir",
|
||||
type=str,
|
||||
default=os.path.join(os.path.dirname(__file__), "sample_data"),
|
||||
help="Path to the data folder in which to save outputs, pdfs should be in /pdfs folder within it.",
|
||||
)
|
||||
parser.add_argument("--force", action="store_true", default=False, help="Force regenerating of output files, even if they already exist")
|
||||
parser.add_argument("--parallel", type=int, default=1, help="Maximum number of concurrent tasks")
|
||||
parser.add_argument(
|
||||
"--remove_text",
|
||||
action="store_true",
|
||||
help="When your PDF gets processed, we will take a screenshot of it first, to erase any text content in it. This would disable document-anchoring for olmocr.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--failfast",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Fail immediately and print full stack trace if any page generation throws an exception",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Mapping of method names to a tuple: (module path, function name)
|
||||
available_methods = {
|
||||
"olmocr_pipeline": ("olmocr.bench.runners.run_olmocr_pipeline", "run_olmocr_pipeline"),
|
||||
"gotocr": ("olmocr.bench.runners.run_gotocr", "run_gotocr"),
|
||||
"nanonetsocr": ("olmocr.bench.runners.run_nanonetsocr", "run_nanonetsocr"),
|
||||
"nanonetsocr_2": ("olmocr.bench.runners.run_nanonetsocr_2", "run_server"),
|
||||
"marker": ("olmocr.bench.runners.run_marker", "run_marker"),
|
||||
"mineru": ("olmocr.bench.runners.run_mineru", "run_mineru"),
|
||||
"chatgpt": ("olmocr.bench.runners.run_chatgpt", "run_chatgpt"),
|
||||
"gemini": ("olmocr.bench.runners.run_gemini", "run_gemini"),
|
||||
"mistral": ("olmocr.bench.runners.run_mistral", "run_mistral"),
|
||||
"docling": ("olmocr.bench.runners.run_docling", "run_docling"),
|
||||
"dotsocr": ("olmocr.bench.runners.run_dotsocr", "run_dotsocr"),
|
||||
"rolmocr": ("olmocr.bench.runners.run_rolmocr", "run_rolmocr"),
|
||||
"paddlepaddle": ("olmocr.bench.runners.run_paddlepaddle", "run_paddlepaddle"),
|
||||
"paddlevl": ("olmocr.bench.runners.run_paddlevl", "run_paddlevl"),
|
||||
"transformers": ("olmocr.bench.runners.run_transformers", "run_transformers"),
|
||||
"server": ("olmocr.bench.runners.run_server", "run_server"),
|
||||
}
|
||||
|
||||
# Build config by importing only requested methods.
|
||||
config = {}
|
||||
for method_arg in args.methods:
|
||||
method_name, extra_kwargs, folder_name = parse_method_arg(method_arg)
|
||||
if method_name not in available_methods:
|
||||
parser.error(f"Unknown method: {method_name}. " f"Available methods: {', '.join(available_methods.keys())}")
|
||||
module_path, function_name = available_methods[method_name]
|
||||
# Dynamically import the module and get the function.
|
||||
module = importlib.import_module(module_path)
|
||||
function = getattr(module, function_name)
|
||||
config[method_name] = {"method": function, "kwargs": extra_kwargs, "folder_name": folder_name}
|
||||
|
||||
data_directory = args.dir
|
||||
pdf_directory = os.path.join(data_directory, "pdfs")
|
||||
|
||||
# Run the async process function with the parallel and failfast arguments
|
||||
asyncio.run(process_pdfs(config, pdf_directory, data_directory, args.repeats, args.remove_text, args.force, args.parallel, args.failfast))
|
||||
@@ -0,0 +1 @@
|
||||
from .render import compare_rendered_equations, render_equation
|
||||
+1
@@ -0,0 +1 @@
|
||||
!function(e,t){"object"==typeof exports&&"object"==typeof module?module.exports=t(require("katex")):"function"==typeof define&&define.amd?define(["katex"],t):"object"==typeof exports?exports.renderMathInElement=t(require("katex")):e.renderMathInElement=t(e.katex)}("undefined"!=typeof self?self:this,(function(e){return function(){"use strict";var t={757:function(t){t.exports=e}},n={};function r(e){var o=n[e];if(void 0!==o)return o.exports;var i=n[e]={exports:{}};return t[e](i,i.exports,r),i.exports}r.n=function(e){var t=e&&e.__esModule?function(){return e.default}:function(){return e};return r.d(t,{a:t}),t},r.d=function(e,t){for(var n in t)r.o(t,n)&&!r.o(e,n)&&Object.defineProperty(e,n,{enumerable:!0,get:t[n]})},r.o=function(e,t){return Object.prototype.hasOwnProperty.call(e,t)};var o={};r.d(o,{default:function(){return p}});var i=r(757),a=r.n(i);const l=function(e,t,n){let r=n,o=0;const i=e.length;for(;r<t.length;){const n=t[r];if(o<=0&&t.slice(r,r+i)===e)return r;"\\"===n?r++:"{"===n?o++:"}"===n&&o--,r++}return-1},s=/^\\begin{/;var d=function(e,t){let n;const r=[],o=new RegExp("("+t.map((e=>e.left.replace(/[-/\\^$*+?.()|[\]{}]/g,"\\$&"))).join("|")+")");for(;n=e.search(o),-1!==n;){n>0&&(r.push({type:"text",data:e.slice(0,n)}),e=e.slice(n));const o=t.findIndex((t=>e.startsWith(t.left)));if(n=l(t[o].right,e,t[o].left.length),-1===n)break;const i=e.slice(0,n+t[o].right.length),a=s.test(i)?i:e.slice(t[o].left.length,n);r.push({type:"math",data:a,rawData:i,display:t[o].display}),e=e.slice(n+t[o].right.length)}return""!==e&&r.push({type:"text",data:e}),r};const c=function(e,t){const n=d(e,t.delimiters);if(1===n.length&&"text"===n[0].type)return null;const r=document.createDocumentFragment();for(let e=0;e<n.length;e++)if("text"===n[e].type)r.appendChild(document.createTextNode(n[e].data));else{const o=document.createElement("span");let i=n[e].data;t.displayMode=n[e].display;try{t.preProcess&&(i=t.preProcess(i)),a().render(i,o,t)}catch(o){if(!(o instanceof a().ParseError))throw o;t.errorCallback("KaTeX auto-render: Failed to parse `"+n[e].data+"` with ",o),r.appendChild(document.createTextNode(n[e].rawData));continue}r.appendChild(o)}return r},f=function(e,t){for(let n=0;n<e.childNodes.length;n++){const r=e.childNodes[n];if(3===r.nodeType){let o=r.textContent,i=r.nextSibling,a=0;for(;i&&i.nodeType===Node.TEXT_NODE;)o+=i.textContent,i=i.nextSibling,a++;const l=c(o,t);if(l){for(let e=0;e<a;e++)r.nextSibling.remove();n+=l.childNodes.length-1,e.replaceChild(l,r)}else n+=a}else if(1===r.nodeType){const e=" "+r.className+" ";-1===t.ignoredTags.indexOf(r.nodeName.toLowerCase())&&t.ignoredClasses.every((t=>-1===e.indexOf(" "+t+" ")))&&f(r,t)}}};var p=function(e,t){if(!e)throw new Error("No element provided to render");const n={};for(const e in t)t.hasOwnProperty(e)&&(n[e]=t[e]);n.delimiters=n.delimiters||[{left:"$$",right:"$$",display:!0},{left:"\\(",right:"\\)",display:!1},{left:"\\begin{equation}",right:"\\end{equation}",display:!0},{left:"\\begin{align}",right:"\\end{align}",display:!0},{left:"\\begin{alignat}",right:"\\end{alignat}",display:!0},{left:"\\begin{gather}",right:"\\end{gather}",display:!0},{left:"\\begin{CD}",right:"\\end{CD}",display:!0},{left:"\\[",right:"\\]",display:!0}],n.ignoredTags=n.ignoredTags||["script","noscript","style","textarea","pre","code","option"],n.ignoredClasses=n.ignoredClasses||[],n.errorCallback=n.errorCallback||console.error,n.macros=n.macros||{},f(e,n)};return o=o.default}()}));
|
||||
Vendored
+1
File diff suppressed because one or more lines are too long
Vendored
+1
File diff suppressed because one or more lines are too long
@@ -0,0 +1,553 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Extract inner-most spans and their bounding boxes, and the MathML output,
|
||||
from rendered LaTeX equations using Playwright and KaTeX.
|
||||
Caching is maintained via a SHA1-based hash stored in a sqlite database.
|
||||
|
||||
Requirements:
|
||||
pip install playwright
|
||||
python -m playwright install chromium
|
||||
|
||||
Place katex.min.css and katex.min.js in the same directory as this script
|
||||
"""
|
||||
|
||||
import atexit
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import re
|
||||
import sqlite3
|
||||
import threading
|
||||
import unittest
|
||||
import weakref
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import List, Optional
|
||||
|
||||
from playwright.sync_api import Error as PlaywrightError
|
||||
from playwright.sync_api import sync_playwright
|
||||
|
||||
# --- New SQLite Cache Implementation ---
|
||||
|
||||
|
||||
class EquationCache:
|
||||
def __init__(self, db_path: Optional[str] = None):
|
||||
if db_path is None:
|
||||
# Use the same cache directory as before
|
||||
cache_dir = pathlib.Path.home() / ".cache" / "olmocr" / "bench" / "equations"
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
db_path = str(cache_dir / "cache.db")
|
||||
self.db_path = db_path
|
||||
self.lock = threading.Lock()
|
||||
self._init_db()
|
||||
|
||||
def _init_db(self):
|
||||
with self.lock:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
c = conn.cursor()
|
||||
# Added an 'error' column to store rendering errors
|
||||
c.execute("""
|
||||
CREATE TABLE IF NOT EXISTS equations (
|
||||
eq_hash TEXT PRIMARY KEY,
|
||||
mathml TEXT,
|
||||
spans TEXT,
|
||||
error TEXT
|
||||
)
|
||||
""")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
def load(self, eq_hash: str) -> Optional["RenderedEquation"]:
|
||||
with self.lock:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
c = conn.cursor()
|
||||
c.execute("SELECT mathml, spans, error FROM equations WHERE eq_hash = ?", (eq_hash,))
|
||||
row = c.fetchone()
|
||||
conn.close()
|
||||
if row:
|
||||
mathml, spans_json, error = row
|
||||
if error:
|
||||
# In error cases, we return an instance with error set and no spans.
|
||||
return RenderedEquation(mathml=mathml, spans=[], error=error)
|
||||
else:
|
||||
spans_data = json.loads(spans_json)
|
||||
spans = [
|
||||
SpanInfo(
|
||||
text=s["text"],
|
||||
bounding_box=BoundingBox(
|
||||
x=s["boundingBox"]["x"],
|
||||
y=s["boundingBox"]["y"],
|
||||
width=s["boundingBox"]["width"],
|
||||
height=s["boundingBox"]["height"],
|
||||
),
|
||||
)
|
||||
for s in spans_data
|
||||
]
|
||||
return RenderedEquation(mathml=mathml, spans=spans)
|
||||
return None
|
||||
|
||||
def save(self, eq_hash: str, rendered_eq: "RenderedEquation"):
|
||||
spans_data = [
|
||||
{
|
||||
"text": span.text,
|
||||
"boundingBox": {
|
||||
"x": span.bounding_box.x,
|
||||
"y": span.bounding_box.y,
|
||||
"width": span.bounding_box.width,
|
||||
"height": span.bounding_box.height,
|
||||
},
|
||||
}
|
||||
for span in rendered_eq.spans
|
||||
]
|
||||
spans_json = json.dumps(spans_data)
|
||||
with self.lock:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
c = conn.cursor()
|
||||
c.execute(
|
||||
"INSERT OR REPLACE INTO equations (eq_hash, mathml, spans, error) VALUES (?, ?, ?, ?)",
|
||||
(eq_hash, rendered_eq.mathml, spans_json, rendered_eq.error),
|
||||
)
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
def clear(self):
|
||||
with self.lock:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
c = conn.cursor()
|
||||
c.execute("DELETE FROM equations")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
|
||||
# Global instance of EquationCache
|
||||
equation_cache = EquationCache()
|
||||
|
||||
# --- End SQLite Cache Implementation ---
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BoundingBox:
|
||||
x: float
|
||||
y: float
|
||||
width: float
|
||||
height: float
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SpanInfo:
|
||||
text: str
|
||||
bounding_box: BoundingBox
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RenderedEquation:
|
||||
mathml: str
|
||||
spans: List[SpanInfo]
|
||||
error: Optional[str] = None # New field to store error messages if rendering fails
|
||||
|
||||
|
||||
def get_equation_hash(equation, bg_color="white", text_color="black", font_size=24):
|
||||
"""
|
||||
Calculate SHA1 hash of the equation string and rendering parameters.
|
||||
"""
|
||||
params_str = f"{equation}|{bg_color}|{text_color}|{font_size}"
|
||||
return hashlib.sha1(params_str.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
# Thread-local storage for browser instances in the executor threads
|
||||
_thread_local = threading.local()
|
||||
|
||||
# Global thread pool executor with a fixed number of threads
|
||||
# Each thread will maintain its own Playwright instance
|
||||
_render_executor = ThreadPoolExecutor(max_workers=8, thread_name_prefix="playwright-render")
|
||||
|
||||
|
||||
def _cleanup_executor():
|
||||
"""Cleanup function to shutdown the executor on exit."""
|
||||
_render_executor.shutdown(wait=False)
|
||||
|
||||
|
||||
# Register cleanup at exit
|
||||
atexit.register(_cleanup_executor)
|
||||
|
||||
|
||||
def _cleanup_playwright(playwright, browser):
|
||||
print("Cleaning up", playwright)
|
||||
try:
|
||||
browser.close()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
playwright.stop()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class _BrowserOwner:
|
||||
def __init__(self):
|
||||
p = sync_playwright().start()
|
||||
b = p.chromium.launch()
|
||||
self.p = p
|
||||
self.browser = b
|
||||
self._closed = False
|
||||
# Important: don't capture `self` or globals in the finalizer
|
||||
self._finalizer = weakref.finalize(self, _cleanup_playwright, p, b)
|
||||
|
||||
def close_now(self):
|
||||
if not self._closed:
|
||||
self._closed = True
|
||||
self._finalizer() # idempotent; runs at most once
|
||||
|
||||
|
||||
def _get_thread_local_browser():
|
||||
"""Get or create a browser instance for the current thread."""
|
||||
owner = getattr(_thread_local, "owner", None)
|
||||
if owner is None:
|
||||
owner = _BrowserOwner()
|
||||
_thread_local.owner = owner
|
||||
return owner
|
||||
|
||||
|
||||
def _render_in_executor(equation, bg_color, text_color, font_size, use_cache, debug_dom, eq_hash):
|
||||
"""
|
||||
Function to be run in the executor thread pool.
|
||||
Each thread maintains its own Playwright instance.
|
||||
"""
|
||||
owner = _get_thread_local_browser()
|
||||
ctx = owner.browser.new_context(viewport={"width": 800, "height": 400})
|
||||
try:
|
||||
return _do_render(ctx, equation, bg_color, text_color, font_size, debug_dom)
|
||||
finally:
|
||||
try:
|
||||
ctx.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _do_render(context, equation, bg_color, text_color, font_size, debug_dom):
|
||||
"""
|
||||
Internal rendering function that uses a provided browser context.
|
||||
"""
|
||||
# Escape the equation for use in a JavaScript string.
|
||||
escaped_equation = json.dumps(equation)
|
||||
|
||||
# Get local paths for KaTeX files.
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
katex_css_path = os.path.join(script_dir, "katex.min.css")
|
||||
katex_js_path = os.path.join(script_dir, "katex.min.js")
|
||||
|
||||
if not os.path.exists(katex_css_path) or not os.path.exists(katex_js_path):
|
||||
raise FileNotFoundError(f"KaTeX files not found. Please ensure katex.min.css and katex.min.js are in {script_dir}")
|
||||
|
||||
# Create a new page.
|
||||
page = context.new_page()
|
||||
|
||||
# Basic HTML structure for rendering.
|
||||
page_html = f"""
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<style>
|
||||
body {{
|
||||
display: flex;
|
||||
justify-content: center;
|
||||
align-items: center;
|
||||
height: 100vh;
|
||||
margin: 0;
|
||||
background-color: {bg_color};
|
||||
color: {text_color};
|
||||
}}
|
||||
#equation-container {{
|
||||
padding: 0;
|
||||
font-size: {font_size}px;
|
||||
}}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div id="equation-container"></div>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
page.set_content(page_html)
|
||||
page.add_style_tag(path=katex_css_path)
|
||||
page.add_script_tag(path=katex_js_path)
|
||||
page.wait_for_load_state("networkidle", timeout=0)
|
||||
|
||||
katex_loaded = page.evaluate("typeof katex !== 'undefined'")
|
||||
if not katex_loaded:
|
||||
page.close()
|
||||
raise RuntimeError("KaTeX library failed to load. Check your katex.min.js file.")
|
||||
|
||||
try:
|
||||
error_message = page.evaluate(f"""
|
||||
() => {{
|
||||
try {{
|
||||
katex.render({escaped_equation}, document.getElementById("equation-container"), {{
|
||||
displayMode: true,
|
||||
throwOnError: true
|
||||
}});
|
||||
return null;
|
||||
}} catch (error) {{
|
||||
console.error("KaTeX error:", error.message);
|
||||
return error.message;
|
||||
}}
|
||||
}}
|
||||
""")
|
||||
except PlaywrightError as ex:
|
||||
print(escaped_equation)
|
||||
error_message = str(ex)
|
||||
page.close()
|
||||
raise
|
||||
|
||||
if error_message:
|
||||
print(f"Error rendering equation: '{equation}'")
|
||||
print(error_message)
|
||||
# Return error result
|
||||
page.close()
|
||||
return RenderedEquation(mathml=error_message, spans=[], error=error_message)
|
||||
|
||||
page.wait_for_selector(".katex", state="attached", timeout=0)
|
||||
|
||||
if debug_dom:
|
||||
katex_dom_html = page.evaluate("""
|
||||
() => {
|
||||
return document.getElementById("equation-container").innerHTML;
|
||||
}
|
||||
""")
|
||||
print("\n===== KaTeX DOM HTML =====")
|
||||
print(katex_dom_html)
|
||||
|
||||
# Extract inner-most spans with non-whitespace text.
|
||||
spans_info = page.evaluate("""
|
||||
() => {
|
||||
const spans = Array.from(document.querySelectorAll('span'));
|
||||
const list = [];
|
||||
spans.forEach(span => {
|
||||
if (span.children.length === 0 && /\\S/.test(span.textContent)) {
|
||||
const rect = span.getBoundingClientRect();
|
||||
list.push({
|
||||
text: span.textContent.trim(),
|
||||
boundingBox: {
|
||||
x: rect.x,
|
||||
y: rect.y,
|
||||
width: rect.width,
|
||||
height: rect.height
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
return list;
|
||||
}
|
||||
""")
|
||||
|
||||
if debug_dom:
|
||||
print("\n===== Extracted Span Information =====")
|
||||
print(spans_info)
|
||||
|
||||
# Extract MathML output (if available) from the KaTeX output.
|
||||
mathml = page.evaluate("""
|
||||
() => {
|
||||
const mathElem = document.querySelector('.katex-mathml math');
|
||||
return mathElem ? mathElem.outerHTML : "";
|
||||
}
|
||||
""")
|
||||
|
||||
page.close()
|
||||
|
||||
rendered_eq = RenderedEquation(
|
||||
mathml=mathml,
|
||||
spans=[
|
||||
SpanInfo(
|
||||
text=s["text"],
|
||||
bounding_box=BoundingBox(
|
||||
x=s["boundingBox"]["x"],
|
||||
y=s["boundingBox"]["y"],
|
||||
width=s["boundingBox"]["width"],
|
||||
height=s["boundingBox"]["height"],
|
||||
),
|
||||
)
|
||||
for s in spans_info
|
||||
],
|
||||
)
|
||||
|
||||
return rendered_eq
|
||||
|
||||
|
||||
def render_equation(
|
||||
equation,
|
||||
bg_color="white",
|
||||
text_color="black",
|
||||
font_size=24,
|
||||
use_cache=True,
|
||||
debug_dom=False,
|
||||
):
|
||||
"""
|
||||
Render a LaTeX equation using Playwright and KaTeX, extract the inner-most span elements
|
||||
along with their bounding boxes, and extract the MathML output generated by KaTeX.
|
||||
|
||||
This function uses a ThreadPoolExecutor with a fixed number of threads to prevent
|
||||
resource leaks from unbounded thread creation.
|
||||
"""
|
||||
# Calculate hash for caching.
|
||||
eq_hash = get_equation_hash(equation, bg_color, text_color, font_size)
|
||||
|
||||
# Try to load from SQLite cache.
|
||||
if use_cache:
|
||||
cached = equation_cache.load(eq_hash)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
# Submit the rendering task to the thread pool executor
|
||||
future = _render_executor.submit(_render_in_executor, equation, bg_color, text_color, font_size, use_cache, debug_dom, eq_hash)
|
||||
|
||||
# Wait for the result
|
||||
rendered_eq = future.result()
|
||||
|
||||
# Save to cache if successful and caching is enabled
|
||||
if use_cache and rendered_eq and not rendered_eq.error:
|
||||
equation_cache.save(eq_hash, rendered_eq)
|
||||
|
||||
return rendered_eq
|
||||
|
||||
|
||||
def compare_rendered_equations(reference: RenderedEquation, hypothesis: RenderedEquation) -> bool:
|
||||
"""
|
||||
Compare two RenderedEquation objects.
|
||||
First, check if the normalized MathML of the hypothesis is contained within that of the reference.
|
||||
If not, perform a neighbor-based matching on the spans.
|
||||
"""
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
def extract_inner(mathml: str) -> str:
|
||||
try:
|
||||
soup = BeautifulSoup(mathml, "xml")
|
||||
semantics = soup.find("semantics")
|
||||
if semantics:
|
||||
inner_parts = [str(child) for child in semantics.contents if getattr(child, "name", None) != "annotation"]
|
||||
return "".join(inner_parts)
|
||||
else:
|
||||
return str(soup)
|
||||
except Exception as e:
|
||||
print("Error parsing MathML with BeautifulSoup:", e)
|
||||
print(mathml)
|
||||
return mathml
|
||||
|
||||
# Filters out all spaces, include unicode zero and partial width spaces
|
||||
def normalize(s: str) -> str:
|
||||
return re.sub(r"[\s\u200a-\u200d\u2060\ufeff]+", "", s, flags=re.UNICODE)
|
||||
|
||||
reference_inner = normalize(extract_inner(reference.mathml))
|
||||
hypothesis_inner = normalize(extract_inner(hypothesis.mathml))
|
||||
if reference_inner in hypothesis_inner:
|
||||
return True
|
||||
|
||||
H, R = reference.spans, hypothesis.spans
|
||||
|
||||
# Filter out any whitespace only spans and remove whitespace from the spans
|
||||
H = [replace(span, text=normalize(span.text)) for span in H if normalize(span.text)]
|
||||
R = [replace(span, text=normalize(span.text)) for span in R if normalize(span.text)]
|
||||
|
||||
def expand_span_info(span_info: SpanInfo) -> list[SpanInfo]:
|
||||
total_elems = len(span_info.text)
|
||||
return [
|
||||
SpanInfo(
|
||||
c,
|
||||
BoundingBox(
|
||||
span_info.bounding_box.x + (span_info.bounding_box.width * index) / total_elems,
|
||||
span_info.bounding_box.y,
|
||||
span_info.bounding_box.width / total_elems,
|
||||
span_info.bounding_box.height,
|
||||
),
|
||||
)
|
||||
for index, c in enumerate(span_info.text)
|
||||
]
|
||||
|
||||
H = [span for sublist in H for span in expand_span_info(sublist)]
|
||||
R = [span for sublist in R for span in expand_span_info(sublist)]
|
||||
|
||||
candidate_map = {}
|
||||
for i, hspan in enumerate(H):
|
||||
candidate_map[i] = [j for j, rsp in enumerate(R) if rsp.text == hspan.text]
|
||||
if not candidate_map[i]:
|
||||
return False
|
||||
|
||||
def compute_neighbors(spans, tol=5):
|
||||
neighbors = {}
|
||||
for i, span in enumerate(spans):
|
||||
cx = span.bounding_box.x + span.bounding_box.width / 2
|
||||
cy = span.bounding_box.y + span.bounding_box.height / 2
|
||||
up = down = left = right = None
|
||||
up_dist = down_dist = left_dist = right_dist = None
|
||||
for j, other in enumerate(spans):
|
||||
if i == j:
|
||||
continue
|
||||
ocx = other.bounding_box.x + other.bounding_box.width / 2
|
||||
ocy = other.bounding_box.y + other.bounding_box.height / 2
|
||||
if ocy < cy and abs(ocx - cx) <= tol:
|
||||
dist = cy - ocy
|
||||
if up is None or dist < up_dist:
|
||||
up = j
|
||||
up_dist = dist
|
||||
if ocy > cy and abs(ocx - cx) <= tol:
|
||||
dist = ocy - cy
|
||||
if down is None or dist < down_dist:
|
||||
down = j
|
||||
down_dist = dist
|
||||
if ocx < cx and abs(ocy - cy) <= tol:
|
||||
dist = cx - ocx
|
||||
if left is None or dist < left_dist:
|
||||
left = j
|
||||
left_dist = dist
|
||||
if ocx > cx and abs(ocy - cy) <= tol:
|
||||
dist = ocx - cx
|
||||
if right is None or dist < right_dist:
|
||||
right = j
|
||||
right_dist = dist
|
||||
neighbors[i] = {"up": up, "down": down, "left": left, "right": right}
|
||||
return neighbors
|
||||
|
||||
hyp_neighbors = compute_neighbors(H)
|
||||
ref_neighbors = compute_neighbors(R)
|
||||
|
||||
n = len(H)
|
||||
used = [False] * len(R)
|
||||
assignment = {}
|
||||
|
||||
def backtrack(i):
|
||||
if i == n:
|
||||
return True
|
||||
for cand in candidate_map[i]:
|
||||
if used[cand]:
|
||||
continue
|
||||
assignment[i] = cand
|
||||
used[cand] = True
|
||||
valid = True
|
||||
for direction in ["up", "down", "left", "right"]:
|
||||
hyp_nb = hyp_neighbors[i].get(direction)
|
||||
ref_nb = ref_neighbors[cand].get(direction)
|
||||
if hyp_nb is not None:
|
||||
expected_text = H[hyp_nb].text
|
||||
if ref_nb is None:
|
||||
valid = False
|
||||
break
|
||||
if hyp_nb in assignment:
|
||||
if assignment[hyp_nb] != ref_nb:
|
||||
valid = False
|
||||
break
|
||||
else:
|
||||
if R[ref_nb].text != expected_text:
|
||||
valid = False
|
||||
break
|
||||
if valid:
|
||||
if backtrack(i + 1):
|
||||
return True
|
||||
used[cand] = False
|
||||
del assignment[i]
|
||||
return False
|
||||
|
||||
return backtrack(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,140 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
|
||||
|
||||
def verify_header_footer_match(
|
||||
pdf_path: str,
|
||||
page_num: int,
|
||||
hea_foo_text: str,
|
||||
model: str,
|
||||
temperature: float = 0.1,
|
||||
target_longest_image_dim: int = 2048,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Verify if a headers and footers matches what appears in a PDF page.
|
||||
|
||||
Args:
|
||||
pdf_path (str): Path to the PDF file
|
||||
page_num (int): Page number to check (1-indexed)
|
||||
model (str): OpenAI model to use
|
||||
temperature (float): Temperature for API call
|
||||
target_longest_image_dim (int): Target dimension for the image
|
||||
|
||||
Returns:
|
||||
Dict with verification result
|
||||
"""
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=target_longest_image_dim)
|
||||
|
||||
if not os.getenv("OPENAI_API_KEY"):
|
||||
raise SystemExit("You must specify an OPENAI_API_KEY environment variable")
|
||||
|
||||
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
|
||||
|
||||
prompt = f"""
|
||||
This is a header and footer verification task.
|
||||
|
||||
I'm showing you a page from a PDF document containing headers and footers text.
|
||||
|
||||
Please verify if the headers or footers are exactly matches the below text.
|
||||
|
||||
{hea_foo_text}
|
||||
|
||||
Respond with a JSON object containing:
|
||||
1. "status": "correct" or "incorrect"
|
||||
2. "confidence": a value between 0 and 1 representing your confidence in the answer
|
||||
3. "explanation": a brief explanation of why you believe the text is correct or incorrect
|
||||
|
||||
Focus specifically on checking if this exact header or footer expression appears in the document.
|
||||
"""
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
# temperature=temperature,
|
||||
response_format={"type": "json_object"},
|
||||
# max_tokens=1000,
|
||||
)
|
||||
raw_response = response.choices[0].message.content
|
||||
result = json.loads(raw_response)
|
||||
|
||||
return {
|
||||
"pdf": pdf_path,
|
||||
"math": hea_foo_text,
|
||||
"status": result.get("status", "unknown"),
|
||||
"confidence": result.get("confidence", 0),
|
||||
"explanation": result.get("explanation", "No explanation provided"),
|
||||
}
|
||||
|
||||
|
||||
def process_jsonl_file(input_jsonl_path: str, output_jsonl_path: str, model: str = "o3-2025-04-16", temperature: float = 0.1) -> None:
|
||||
"""
|
||||
Process a JSONL file containing math expressions to verify.
|
||||
|
||||
Args:
|
||||
input_jsonl_path (str): Path to input JSONL file
|
||||
output_jsonl_path (str): Path to output JSONL file
|
||||
model (str): OpenAI model to use
|
||||
temperature (float): Temperature for API call
|
||||
"""
|
||||
processed_count = 0
|
||||
|
||||
with open(output_jsonl_path, "w") as out_file:
|
||||
with open(input_jsonl_path, "r") as in_file:
|
||||
for line_num, line in enumerate(in_file, 1):
|
||||
try:
|
||||
entry = json.loads(line.strip())
|
||||
|
||||
pdf_path = entry.get("pdf")
|
||||
page_num = entry.get("page", 1)
|
||||
text_expr = entry.get("text")
|
||||
|
||||
if not all([pdf_path, text_expr]):
|
||||
print(f"Line {line_num}: Skipping entry due to missing required fields")
|
||||
continue
|
||||
|
||||
print(f"Line {line_num}: Processing: {pdf_path}, page {page_num}")
|
||||
|
||||
try:
|
||||
result = verify_header_footer_match(pdf_path=pdf_path, page_num=page_num, hea_foo_text=text_expr, model=model, temperature=temperature)
|
||||
out_file.write(json.dumps(result) + "\n")
|
||||
processed_count += 1
|
||||
except Exception as e:
|
||||
print(f"Line {line_num}: Error processing {pdf_path}: {str(e)}")
|
||||
error_result = {"pdf": pdf_path, "text": text_expr, "status": "error", "explanation": str(e)}
|
||||
out_file.write(json.dumps(error_result) + "\n")
|
||||
processed_count += 1
|
||||
|
||||
except json.JSONDecodeError:
|
||||
print(f"Line {line_num}: Invalid JSON, skipping")
|
||||
|
||||
print(f"Processed {processed_count} entries. Results saved to {output_jsonl_path}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Verify headers footers expressions in PDFs")
|
||||
parser.add_argument("input_jsonl", help="Path to input JSONL file")
|
||||
parser.add_argument("output_jsonl", help="Path to output JSONL file")
|
||||
parser.add_argument("--model", default="o3-2025-04-16", help="OpenAI model to use")
|
||||
parser.add_argument("--temperature", type=float, default=0.1, help="Temperature for API call")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
process_jsonl_file(input_jsonl_path=args.input_jsonl, output_jsonl_path=args.output_jsonl, model=args.model, temperature=args.temperature)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+159
@@ -0,0 +1,159 @@
|
||||
#!/usr/bin/env python3
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Any, Dict
|
||||
|
||||
import openai
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
|
||||
|
||||
def process_test_case(case: Dict[str, Any], client, pdf_dir: str, model: str = "gpt-4o") -> Dict[str, Any]:
|
||||
"""
|
||||
Send a request to GPT-4 asking if the before and after text appear in the same region.
|
||||
Include the PDF image in the prompt.
|
||||
|
||||
Args:
|
||||
case: A test case from the JSONL file
|
||||
client: The OpenAI client
|
||||
pdf_dir: Directory containing PDF files
|
||||
model: The model to use
|
||||
|
||||
Returns:
|
||||
The original case with the added response field
|
||||
"""
|
||||
before_text = case["before"]
|
||||
after_text = case["after"]
|
||||
pdf_path = os.path.join(pdf_dir, case["pdf"])
|
||||
page_num = case["page"]
|
||||
|
||||
try:
|
||||
# Render the PDF page to a base64-encoded PNG image
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num)
|
||||
|
||||
# Create messages with both text and image
|
||||
messages = [
|
||||
{"role": "system", "content": "You are an AI assistant analyzing text from PDFs."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
f"Does the text in the 'before' field and the 'after' field appear in the same region of the page? "
|
||||
f"Look at the PDF image and determine if these texts are located near each other or in completely "
|
||||
f"different parts of the page. Different regions could be the captions for different images, or inside of different insets or tables. However, appearing the same column of text, or in the naturally flowing next column of text is close enough.\n\n"
|
||||
f"Before: {before_text}\n\n"
|
||||
f"After: {after_text}\n\n"
|
||||
f"Respond with 'YES' if they appear in the same region or column, and 'NO' if they appear in "
|
||||
f"different regions. Then explain your reasoning in 1-2 sentences."
|
||||
),
|
||||
},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
# Call the API
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.0,
|
||||
max_tokens=300,
|
||||
)
|
||||
|
||||
# Add GPT-4's response to the case
|
||||
case_with_response = case.copy()
|
||||
case_with_response["gpt4_response"] = response.choices[0].message.content
|
||||
return case_with_response
|
||||
except Exception as e:
|
||||
# In case of error, return the original case with an error message
|
||||
case_with_response = case.copy()
|
||||
case_with_response["gpt4_response"] = f"ERROR: {str(e)}"
|
||||
# Print the error for debugging
|
||||
print(f"Error processing {case.get('id', 'unknown')}: {str(e)}")
|
||||
return case_with_response
|
||||
|
||||
|
||||
def process_jsonl_file(input_file: str, output_file: str, api_key: str, pdf_dir: str, num_workers: int = 8, model: str = "gpt-4o") -> None:
|
||||
"""
|
||||
Process each line in the JSONL file by sending requests to GPT-4 in parallel.
|
||||
|
||||
Args:
|
||||
input_file: Path to the input JSONL file
|
||||
output_file: Path to write the output JSONL file with responses
|
||||
api_key: OpenAI API key
|
||||
pdf_dir: Directory containing PDF files
|
||||
num_workers: Number of parallel workers
|
||||
model: The model to use
|
||||
"""
|
||||
# Read all test cases from the input file
|
||||
with open(input_file, "r") as f:
|
||||
lines = f.readlines()
|
||||
|
||||
# Parse each line to get test cases
|
||||
test_cases = []
|
||||
for line in lines:
|
||||
if line.strip():
|
||||
test_cases.append(json.loads(line))
|
||||
|
||||
# Initialize OpenAI client
|
||||
client = openai.OpenAI(api_key=api_key)
|
||||
|
||||
# Process test cases in parallel
|
||||
results = []
|
||||
with ThreadPoolExecutor(max_workers=num_workers) as executor:
|
||||
# Submit all tasks
|
||||
future_to_case = {executor.submit(process_test_case, case, client, pdf_dir, model): case for case in test_cases}
|
||||
|
||||
# Process results as they complete
|
||||
for future in tqdm(as_completed(future_to_case), total=len(test_cases), desc="Processing test cases"):
|
||||
try:
|
||||
result = future.result()
|
||||
results.append(result)
|
||||
except Exception as e:
|
||||
case = future_to_case[future]
|
||||
print(f"Error processing case {case.get('id', 'unknown')}: {str(e)}")
|
||||
# Add failed case with error message
|
||||
case["gpt4_response"] = f"PROCESSING_ERROR: {str(e)}"
|
||||
results.append(case)
|
||||
|
||||
# Filter for cases where GPT-4 responded with "NO"
|
||||
no_responses = [result for result in results if "gpt4_response" in result and result["gpt4_response"].startswith("NO")]
|
||||
|
||||
# Write filtered results to output file
|
||||
with open(output_file, "w") as f:
|
||||
for result in no_responses:
|
||||
f.write(json.dumps(result) + "\n")
|
||||
|
||||
print(f"Processed {len(results)} test cases. Found {len(no_responses)} cases with 'NO' responses. Results written to {output_file}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Process multi_column.jsonl with GPT-4 to check text regions")
|
||||
parser.add_argument("--input", default="/home/ubuntu/olmocr/olmOCR-bench/bench_data/multi_column.jsonl", help="Path to input JSONL file")
|
||||
parser.add_argument("--output", default="/home/ubuntu/olmocr/olmOCR-bench/bench_data/multi_column_gpt4_regions.jsonl", help="Path to output JSONL file")
|
||||
parser.add_argument("--pdf-dir", default="/home/ubuntu/olmocr/olmOCR-bench/bench_data/pdfs", help="Directory containing the PDF files")
|
||||
parser.add_argument("--workers", type=int, default=8, help="Number of parallel workers")
|
||||
parser.add_argument("--model", default="gpt-4.1", help="OpenAI model to use")
|
||||
parser.add_argument("--api-key", help="OpenAI API key (if not provided, uses OPENAI_API_KEY env var)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key from arguments or environment variable
|
||||
api_key = args.api_key or os.environ.get("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("OpenAI API key must be provided either via --api-key or OPENAI_API_KEY environment variable")
|
||||
|
||||
# Verify that the PDF directory exists
|
||||
if not os.path.isdir(args.pdf_dir):
|
||||
raise ValueError(f"PDF directory {args.pdf_dir} does not exist")
|
||||
|
||||
process_jsonl_file(input_file=args.input, output_file=args.output, api_key=api_key, pdf_dir=args.pdf_dir, num_workers=args.workers, model=args.model)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,143 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
|
||||
|
||||
def verify_latex_match(
|
||||
pdf_path: str,
|
||||
page_num: int,
|
||||
latex_expression: str,
|
||||
model: str = "gpt-4o-2024-08-06",
|
||||
temperature: float = 0.1,
|
||||
target_longest_image_dim: int = 2048,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Verify if a LaTeX math expression matches what appears in a PDF page.
|
||||
|
||||
Args:
|
||||
pdf_path (str): Path to the PDF file
|
||||
page_num (int): Page number to check (1-indexed)
|
||||
latex_expression (str): LaTeX expression to verify
|
||||
model (str): OpenAI model to use
|
||||
temperature (float): Temperature for API call
|
||||
target_longest_image_dim (int): Target dimension for the image
|
||||
|
||||
Returns:
|
||||
Dict with verification result
|
||||
"""
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=target_longest_image_dim)
|
||||
|
||||
if not os.getenv("OPENAI_API_KEY"):
|
||||
raise SystemExit("You must specify an OPENAI_API_KEY environment variable")
|
||||
|
||||
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
|
||||
|
||||
prompt = f"""
|
||||
This is a mathematical expression verification task.
|
||||
|
||||
I'm showing you a page from a PDF document containing mathematical expressions.
|
||||
|
||||
Please verify if the following LaTeX expression:
|
||||
|
||||
{latex_expression}
|
||||
|
||||
appears correctly in the document.
|
||||
|
||||
Respond with a JSON object containing:
|
||||
1. "status": "correct" or "incorrect"
|
||||
2. "confidence": a value between 0 and 1 representing your confidence in the answer
|
||||
3. "explanation": a brief explanation of why you believe the expression is correct or incorrect
|
||||
|
||||
Focus specifically on checking if this exact mathematical expression appears in the document.
|
||||
"""
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
# temperature=temperature,
|
||||
response_format={"type": "json_object"},
|
||||
# max_tokens=1000,
|
||||
)
|
||||
raw_response = response.choices[0].message.content
|
||||
result = json.loads(raw_response)
|
||||
|
||||
return {
|
||||
"pdf": pdf_path,
|
||||
"math": latex_expression,
|
||||
"status": result.get("status", "unknown"),
|
||||
"confidence": result.get("confidence", 0),
|
||||
"explanation": result.get("explanation", "No explanation provided"),
|
||||
}
|
||||
|
||||
|
||||
def process_jsonl_file(input_jsonl_path: str, output_jsonl_path: str, model: str = "o4-mini-2025-04-16", temperature: float = 0.1) -> None:
|
||||
"""
|
||||
Process a JSONL file containing math expressions to verify.
|
||||
|
||||
Args:
|
||||
input_jsonl_path (str): Path to input JSONL file
|
||||
output_jsonl_path (str): Path to output JSONL file
|
||||
model (str): OpenAI model to use
|
||||
temperature (float): Temperature for API call
|
||||
"""
|
||||
processed_count = 0
|
||||
|
||||
with open(output_jsonl_path, "w") as out_file:
|
||||
with open(input_jsonl_path, "r") as in_file:
|
||||
for line_num, line in enumerate(in_file, 1):
|
||||
try:
|
||||
entry = json.loads(line.strip())
|
||||
|
||||
pdf_path = entry.get("pdf")
|
||||
page_num = entry.get("page", 1)
|
||||
math_expr = entry.get("math")
|
||||
|
||||
if not all([pdf_path, math_expr]):
|
||||
print(f"Line {line_num}: Skipping entry due to missing required fields")
|
||||
continue
|
||||
|
||||
print(f"Line {line_num}: Processing: {pdf_path}, page {page_num}")
|
||||
|
||||
try:
|
||||
result = verify_latex_match(pdf_path=pdf_path, page_num=page_num, latex_expression=math_expr, model=model, temperature=temperature)
|
||||
out_file.write(json.dumps(result) + "\n")
|
||||
processed_count += 1
|
||||
except Exception as e:
|
||||
print(f"Line {line_num}: Error processing {pdf_path}: {str(e)}")
|
||||
error_result = {"pdf": pdf_path, "math": math_expr, "status": "error", "explanation": str(e)}
|
||||
out_file.write(json.dumps(error_result) + "\n")
|
||||
processed_count += 1
|
||||
|
||||
except json.JSONDecodeError:
|
||||
print(f"Line {line_num}: Invalid JSON, skipping")
|
||||
|
||||
print(f"Processed {processed_count} entries. Results saved to {output_jsonl_path}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Verify LaTeX math expressions in PDFs")
|
||||
parser.add_argument("input_jsonl", help="Path to input JSONL file")
|
||||
parser.add_argument("output_jsonl", help="Path to output JSONL file")
|
||||
parser.add_argument("--model", default="o4-mini-2025-04-16", help="OpenAI model to use")
|
||||
parser.add_argument("--temperature", type=float, default=0.1, help="Temperature for API call")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
process_jsonl_file(input_jsonl_path=args.input_jsonl, output_jsonl_path=args.output_jsonl, model=args.model, temperature=args.temperature)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,152 @@
|
||||
#!/usr/bin/env python
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from collections import defaultdict
|
||||
|
||||
from pypdf import PdfReader, PdfWriter
|
||||
|
||||
|
||||
def get_pdf_page_refs(dataset_jsonl):
|
||||
"""
|
||||
Parse dataset.jsonl to extract all PDF page references.
|
||||
Returns a dict mapping (pdf_name, page_num) to a list of test IDs referencing that combination.
|
||||
"""
|
||||
pdf_page_tests = defaultdict(list)
|
||||
|
||||
with open(dataset_jsonl, "r") as f:
|
||||
for line in f:
|
||||
if not line.strip():
|
||||
continue
|
||||
|
||||
try:
|
||||
test = json.loads(line)
|
||||
pdf_name = test.get("pdf")
|
||||
page_num = test.get("page")
|
||||
test_id = test.get("id")
|
||||
|
||||
if pdf_name and page_num and test_id:
|
||||
pdf_page_tests[(pdf_name, page_num)].append(test_id)
|
||||
except json.JSONDecodeError:
|
||||
print(f"Warning: Could not parse line: {line}")
|
||||
continue
|
||||
|
||||
return pdf_page_tests
|
||||
|
||||
|
||||
def extract_single_page_pdfs(source_pdf_dir, target_pdf_dir, pdf_page_tests):
|
||||
"""
|
||||
Extract single page PDFs for each referenced (pdf_name, page_num) combination.
|
||||
"""
|
||||
os.makedirs(target_pdf_dir, exist_ok=True)
|
||||
|
||||
# Track all PDFs we need to process
|
||||
processed_pairs = set()
|
||||
|
||||
for pdf_name, page_num in pdf_page_tests.keys():
|
||||
if (pdf_name, page_num) in processed_pairs:
|
||||
continue
|
||||
|
||||
source_pdf_path = os.path.join(source_pdf_dir, pdf_name)
|
||||
if not os.path.exists(source_pdf_path):
|
||||
print(f"Warning: Source PDF not found: {source_pdf_path}")
|
||||
continue
|
||||
|
||||
try:
|
||||
# Create a new single-page PDF
|
||||
reader = PdfReader(source_pdf_path)
|
||||
writer = PdfWriter()
|
||||
|
||||
# PDF pages are 0-indexed, but our references are 1-indexed
|
||||
zero_indexed_page = page_num - 1
|
||||
|
||||
if zero_indexed_page >= len(reader.pages) or zero_indexed_page < 0:
|
||||
print(f"Warning: Page {page_num} out of range for {pdf_name}")
|
||||
continue
|
||||
|
||||
# Add the specified page to the writer
|
||||
writer.add_page(reader.pages[zero_indexed_page])
|
||||
|
||||
# Create output filename
|
||||
# Remove .pdf extension if it exists
|
||||
base_name = pdf_name.rsplit(".", 1)[0] if pdf_name.lower().endswith(".pdf") else pdf_name
|
||||
output_filename = f"{base_name}_pg{page_num}.pdf"
|
||||
output_path = os.path.join(target_pdf_dir, output_filename)
|
||||
|
||||
# Write the new PDF
|
||||
with open(output_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
processed_pairs.add((pdf_name, page_num))
|
||||
print(f"Created single-page PDF: {output_path}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_name} page {page_num}: {str(e)}")
|
||||
|
||||
return processed_pairs
|
||||
|
||||
|
||||
def reorganize_test_outputs(source_data_dir, target_data_dir, processed_pairs):
|
||||
"""
|
||||
Copy and reorganize test outputs matching the processed PDF page combinations.
|
||||
"""
|
||||
# Create a dataset.jsonl with only the tests for pages we're keeping
|
||||
source_dataset = os.path.join(source_data_dir, "dataset.jsonl")
|
||||
target_dataset = os.path.join(target_data_dir, "dataset.jsonl")
|
||||
|
||||
# Only copy tests for PDFs we processed
|
||||
if os.path.exists(source_dataset):
|
||||
with open(source_dataset, "r") as source_f, open(target_dataset, "w") as target_f:
|
||||
for line in source_f:
|
||||
if not line.strip():
|
||||
continue
|
||||
|
||||
try:
|
||||
test = json.loads(line)
|
||||
pdf_name = test.get("pdf")
|
||||
page_num = test.get("page")
|
||||
|
||||
# Update the PDF name in the test to reflect our new naming convention
|
||||
if (pdf_name, page_num) in processed_pairs:
|
||||
base_name = pdf_name.rsplit(".", 1)[0] if pdf_name.lower().endswith(".pdf") else pdf_name
|
||||
test["pdf"] = f"{base_name}_pg{page_num}.pdf"
|
||||
# Since we've created single-page PDFs, update the page to 1
|
||||
test["page"] = 1
|
||||
target_f.write(json.dumps(test) + "\n")
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Extract single-page PDFs and reorganize test data")
|
||||
parser.add_argument("--source_dir", type=str, required=True, help="Source directory containing the original sample_data structure")
|
||||
parser.add_argument("--target_dir", type=str, required=True, help="Target directory to create the new data structure")
|
||||
args = parser.parse_args()
|
||||
|
||||
source_data_dir = args.source_dir
|
||||
target_data_dir = args.target_dir
|
||||
|
||||
# Create directory structure
|
||||
os.makedirs(target_data_dir, exist_ok=True)
|
||||
target_pdf_dir = os.path.join(target_data_dir, "pdfs")
|
||||
|
||||
# Get paths
|
||||
source_dataset_path = os.path.join(source_data_dir, "dataset.jsonl")
|
||||
source_pdf_dir = os.path.join(source_data_dir, "pdfs")
|
||||
|
||||
# Extract PDF page references from dataset
|
||||
pdf_page_tests = get_pdf_page_refs(source_dataset_path)
|
||||
print(f"Found {len(pdf_page_tests)} unique PDF page combinations referenced in tests")
|
||||
|
||||
# Extract single-page PDFs
|
||||
processed_pairs = extract_single_page_pdfs(source_pdf_dir, target_pdf_dir, pdf_page_tests)
|
||||
print(f"Processed {len(processed_pairs)} unique PDF pages")
|
||||
|
||||
# Reorganize test outputs
|
||||
reorganize_test_outputs(source_data_dir, target_data_dir, processed_pairs)
|
||||
|
||||
print(f"Data extraction complete. New structure created in {target_data_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,64 @@
|
||||
# Rewrites all URLs in a dataset.jsonl file using a sql lite database lookup
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import sqlite3
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def parse_pdf_hash(pretty_pdf_path: str) -> Optional[str]:
|
||||
pattern = r"s3://ai2-s2-pdfs/([a-f0-9]{4})/([a-f0-9]+)\.pdf"
|
||||
match = re.match(pattern, pretty_pdf_path)
|
||||
if match:
|
||||
return match.group(1) + match.group(2)
|
||||
return None
|
||||
|
||||
|
||||
def get_uri_from_db(db_path: str, pdf_hash: str) -> Optional[str]:
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT uri FROM pdf_mapping WHERE pdf_hash = ?", (pdf_hash,))
|
||||
result = cursor.fetchone()
|
||||
conn.close()
|
||||
return result[0] if result else None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Rewrites all URLs in a dataset.jsonl file using a sql lite database lookup")
|
||||
parser.add_argument("jsonl", type=str, help="JSONL file containing s3 paths")
|
||||
parser.add_argument("--db", type=str, required=True, help="Path to sqlite database mapping internal s3 urls to external ones")
|
||||
parser.add_argument("--force", action="store_true", help="Path to sqlite database mapping internal s3 urls to external ones")
|
||||
args = parser.parse_args()
|
||||
|
||||
data = []
|
||||
skipped = 0
|
||||
|
||||
with open(args.jsonl, "r") as inpf:
|
||||
for row in inpf:
|
||||
if len(row.strip()) > 0:
|
||||
j = json.loads(row)
|
||||
|
||||
assert j["url"]
|
||||
hash = parse_pdf_hash(j["url"])
|
||||
if hash:
|
||||
url = get_uri_from_db(args.db, hash)
|
||||
|
||||
if url:
|
||||
j["url"] = url
|
||||
data.append(j)
|
||||
else:
|
||||
skipped += 1
|
||||
else:
|
||||
data.append(j)
|
||||
|
||||
print(data)
|
||||
|
||||
print(f"{skipped} entries were skipped!")
|
||||
|
||||
if not args.force:
|
||||
print("Now run with --force to write data")
|
||||
quit()
|
||||
|
||||
with open(args.jsonl, "w") as inpf:
|
||||
for row in data:
|
||||
print(json.dumps(row), file=inpf)
|
||||
@@ -0,0 +1,220 @@
|
||||
#!/usr/bin/env python
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
def get_rejected_tests(dataset_jsonl):
|
||||
"""
|
||||
Parse dataset.jsonl to identify rejected tests.
|
||||
Returns:
|
||||
- rejected_tests: Set of test IDs that were marked as rejected
|
||||
- pdf_tests: Dict mapping PDF filenames to sets of test IDs
|
||||
- test_pdf_map: Dict mapping test IDs to their PDF filenames
|
||||
"""
|
||||
rejected_tests = set()
|
||||
pdf_tests = defaultdict(set)
|
||||
test_pdf_map = {}
|
||||
|
||||
try:
|
||||
with open(dataset_jsonl, "r") as f:
|
||||
for line in f:
|
||||
if not line.strip():
|
||||
continue
|
||||
|
||||
try:
|
||||
test = json.loads(line)
|
||||
test_id = test.get("id")
|
||||
pdf_name = test.get("pdf")
|
||||
|
||||
# Store the test in our mapping
|
||||
if test_id and pdf_name:
|
||||
pdf_tests[pdf_name].add(test_id)
|
||||
test_pdf_map[test_id] = pdf_name
|
||||
|
||||
# Check if the test is marked as rejected
|
||||
if test.get("checked", None) == "rejected":
|
||||
rejected_tests.add(test_id)
|
||||
|
||||
except json.JSONDecodeError:
|
||||
print(f"Warning: Could not parse line: {line}")
|
||||
continue
|
||||
|
||||
except FileNotFoundError:
|
||||
print(f"Error: Dataset file {dataset_jsonl} not found.")
|
||||
sys.exit(1)
|
||||
|
||||
return rejected_tests, pdf_tests, test_pdf_map
|
||||
|
||||
|
||||
def update_dataset(dataset_jsonl, rejected_tests, dry_run=True):
|
||||
"""
|
||||
Create a new dataset.jsonl without the rejected tests.
|
||||
"""
|
||||
temp_file = dataset_jsonl + ".temp"
|
||||
removed_count = 0
|
||||
|
||||
try:
|
||||
with open(dataset_jsonl, "r") as source, open(temp_file, "w") as target:
|
||||
for line in source:
|
||||
if not line.strip():
|
||||
continue
|
||||
|
||||
try:
|
||||
test = json.loads(line)
|
||||
test_id = test.get("id")
|
||||
|
||||
if test_id in rejected_tests:
|
||||
removed_count += 1
|
||||
else:
|
||||
target.write(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
except FileNotFoundError:
|
||||
print(f"Error: Dataset file {dataset_jsonl} not found.")
|
||||
sys.exit(1)
|
||||
|
||||
if not dry_run:
|
||||
os.replace(temp_file, dataset_jsonl)
|
||||
else:
|
||||
os.remove(temp_file)
|
||||
|
||||
return removed_count
|
||||
|
||||
|
||||
def find_orphaned_pdfs(pdf_dir, pdf_tests, rejected_tests):
|
||||
"""
|
||||
Find PDF files that have all their tests rejected.
|
||||
"""
|
||||
orphaned_pdfs = []
|
||||
|
||||
for pdf_name, tests in pdf_tests.items():
|
||||
# Check if all tests for this PDF are in the rejected list
|
||||
if tests and all(test_id in rejected_tests for test_id in tests):
|
||||
pdf_path = os.path.join(pdf_dir, pdf_name)
|
||||
if os.path.exists(pdf_path):
|
||||
orphaned_pdfs.append(pdf_path)
|
||||
|
||||
return orphaned_pdfs
|
||||
|
||||
|
||||
def find_unreferenced_pdfs(pdf_dir, pdf_tests):
|
||||
"""
|
||||
Find PDF files in the pdf_dir that are not referenced by any test.
|
||||
"""
|
||||
unreferenced_pdfs = []
|
||||
# List all PDFs in the directory (recursively)
|
||||
for pdf_path in glob.glob(os.path.join(pdf_dir, "**", "*.pdf"), recursive=True):
|
||||
# Get the relative path of the PDF from pdf_dir
|
||||
pdf_name = os.path.relpath(pdf_path, pdf_dir)
|
||||
if pdf_name not in pdf_tests:
|
||||
unreferenced_pdfs.append(pdf_path)
|
||||
return unreferenced_pdfs
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Delete rejected tests from dataset and orphaned/unreferenced PDFs")
|
||||
parser.add_argument("--data_dir", type=str, required=True, help="Directory containing dataset.jsonl files and the pdfs/ folder")
|
||||
parser.add_argument("--force", action="store_true", help="Perform actual deletion without confirmation")
|
||||
args = parser.parse_args()
|
||||
|
||||
data_dir = args.data_dir
|
||||
dry_run = not args.force
|
||||
|
||||
# Verify pdfs directory exists
|
||||
pdf_dir = os.path.join(data_dir, "pdfs")
|
||||
if not os.path.exists(pdf_dir):
|
||||
print(f"Error: pdfs/ directory not found in {data_dir}")
|
||||
sys.exit(1)
|
||||
|
||||
# Find all JSONL dataset files in the data_dir
|
||||
dataset_files = glob.glob(os.path.join(data_dir, "*.jsonl"))
|
||||
if not dataset_files:
|
||||
print("No JSONL dataset files found.")
|
||||
sys.exit(0)
|
||||
|
||||
# Global aggregation over all dataset files
|
||||
global_rejected_tests = set()
|
||||
global_pdf_tests = defaultdict(set)
|
||||
global_test_pdf_map = {}
|
||||
|
||||
for dataset_file in dataset_files:
|
||||
rejected_tests, pdf_tests, test_pdf_map = get_rejected_tests(dataset_file)
|
||||
global_rejected_tests |= rejected_tests
|
||||
for pdf_name, test_ids in pdf_tests.items():
|
||||
global_pdf_tests[pdf_name].update(test_ids)
|
||||
global_test_pdf_map.update(test_pdf_map)
|
||||
|
||||
total_tests = sum(len(test_ids) for test_ids in global_pdf_tests.values())
|
||||
|
||||
# Compute orphaned and unreferenced PDFs using global mapping
|
||||
orphaned_pdfs = find_orphaned_pdfs(pdf_dir, global_pdf_tests, global_rejected_tests)
|
||||
unreferenced_pdfs = find_unreferenced_pdfs(pdf_dir, global_pdf_tests)
|
||||
|
||||
# Print summary (global)
|
||||
print("\n===== DELETION SUMMARY =====")
|
||||
print(f"Mode: {'DRY RUN (no changes will be made)' if dry_run else 'FORCE (changes will be applied)'}")
|
||||
print(f"Total tests: {total_tests}")
|
||||
print(f"Tests marked as rejected: {len(global_rejected_tests)}")
|
||||
print(f"PDF files with all tests rejected: {len(orphaned_pdfs)}")
|
||||
print(f"PDF files not referenced by any tests: {len(unreferenced_pdfs)}")
|
||||
|
||||
if global_rejected_tests:
|
||||
print("\nRejected tests:")
|
||||
for test_id in sorted(global_rejected_tests):
|
||||
print(f" - {test_id} (from {global_test_pdf_map.get(test_id, 'unknown')})")
|
||||
|
||||
if orphaned_pdfs:
|
||||
print("\nPDF files to be deleted (all tests rejected):")
|
||||
for pdf_path in sorted(orphaned_pdfs):
|
||||
print(f" - {os.path.basename(pdf_path)}")
|
||||
|
||||
if unreferenced_pdfs:
|
||||
print("\nPDF files to be deleted (unreferenced by any tests):")
|
||||
for pdf_path in sorted(unreferenced_pdfs):
|
||||
print(f" - {os.path.basename(pdf_path)}")
|
||||
|
||||
# If dry run, exit here
|
||||
if dry_run and (global_rejected_tests or orphaned_pdfs or unreferenced_pdfs):
|
||||
print("\nThis is a dry run. No changes have been made.")
|
||||
print("To perform the actual deletion, run the script with the --force flag.")
|
||||
return
|
||||
|
||||
# Confirm before deletion if there are items to delete
|
||||
if global_rejected_tests or orphaned_pdfs or unreferenced_pdfs:
|
||||
confirm = input("\nDo you want to proceed with deletion? (y/N): ")
|
||||
if confirm.lower() not in ("y", "yes"):
|
||||
print("Deletion cancelled.")
|
||||
return
|
||||
|
||||
# Update each dataset file by removing rejected tests
|
||||
for dataset_file in dataset_files:
|
||||
removed_count = update_dataset(dataset_file, global_rejected_tests, dry_run=False)
|
||||
print(f"Removed {removed_count} rejected tests from {os.path.basename(dataset_file)}")
|
||||
|
||||
# Delete orphaned PDFs
|
||||
for pdf_path in orphaned_pdfs:
|
||||
try:
|
||||
os.remove(pdf_path)
|
||||
print(f"Deleted orphaned PDF: {os.path.basename(pdf_path)}")
|
||||
except OSError as e:
|
||||
print(f"Error deleting {os.path.basename(pdf_path)}: {e}")
|
||||
|
||||
# Delete unreferenced PDFs
|
||||
for pdf_path in unreferenced_pdfs:
|
||||
try:
|
||||
os.remove(pdf_path)
|
||||
print(f"Deleted unreferenced PDF: {os.path.basename(pdf_path)}")
|
||||
except OSError as e:
|
||||
print(f"Error deleting {os.path.basename(pdf_path)}: {e}")
|
||||
|
||||
print("\nDeletion completed successfully.")
|
||||
else:
|
||||
print("\nNo rejected tests, orphaned PDFs, or unreferenced PDFs found. Nothing to delete.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,120 @@
|
||||
# This script goes to
|
||||
# https://arxiv.org/list/math/recent?skip=0&show=2000
|
||||
# and downloads all the source PDFs, as well as latex equivalents, and puts them together into
|
||||
# Searching for:
|
||||
# <a href="/pdf/2503.08675" title="Download PDF" id="pdf-2503.08675" aria-labelledby="pdf-2503.08675">pdf</a>
|
||||
# a math_data folder
|
||||
#!/usr/bin/env python3
|
||||
import argparse
|
||||
import io
|
||||
import os
|
||||
import re
|
||||
import tarfile
|
||||
import time
|
||||
|
||||
import requests
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def download_and_extract_source(paper_id, data_dir):
|
||||
source_url = f"https://export.arxiv.org/src/{paper_id}"
|
||||
print(f"Downloading source for {paper_id} from {source_url}...")
|
||||
response = requests.get(source_url)
|
||||
if response.status_code != 200:
|
||||
print(f"Error downloading source for {paper_id}: HTTP {response.status_code}")
|
||||
return False
|
||||
|
||||
# Try to open as a tar archive.
|
||||
try:
|
||||
file_obj = io.BytesIO(response.content)
|
||||
with tarfile.open(fileobj=file_obj, mode="r:*") as tar:
|
||||
# Filter for regular .tex files.
|
||||
members = [m for m in tar.getmembers() if m.isfile() and m.name.endswith(".tex")]
|
||||
print("Found TeX files:", [m.name for m in members])
|
||||
if len(members) == 1:
|
||||
member = members[0]
|
||||
extracted = tar.extractfile(member)
|
||||
if extracted is None:
|
||||
print(f"Error extracting {paper_id}: Could not read the file from the archive.")
|
||||
return False
|
||||
content = extracted.read()
|
||||
out_path = os.path.join(data_dir, f"{paper_id}.tex")
|
||||
with open(out_path, "wb") as f:
|
||||
f.write(content)
|
||||
print(f"Saved tex source for {paper_id} as {out_path}")
|
||||
return True
|
||||
else:
|
||||
print(f"Error: {paper_id} contains multiple .tex files or none. Skipping extraction.")
|
||||
return False
|
||||
except tarfile.ReadError:
|
||||
# Not a tar archive; assume it's a single file.
|
||||
out_path = os.path.join(data_dir, f"{paper_id}.tex")
|
||||
with open(out_path, "wb") as f:
|
||||
f.write(response.content)
|
||||
print(f"Saved non-archive tex source for {paper_id} as {out_path}")
|
||||
return True
|
||||
|
||||
|
||||
def download_pdf(paper_id, data_dir):
|
||||
pdf_url = f"https://export.arxiv.org/pdf/{paper_id}.pdf"
|
||||
print(f"Downloading PDF for {paper_id} from {pdf_url}...")
|
||||
response = requests.get(pdf_url)
|
||||
if response.status_code != 200:
|
||||
print(f"Error downloading PDF for {paper_id}: HTTP {response.status_code}")
|
||||
return False
|
||||
out_path = os.path.join(data_dir, f"{paper_id}.pdf")
|
||||
with open(out_path, "wb") as f:
|
||||
f.write(response.content)
|
||||
print(f"Saved PDF for {paper_id} as {out_path}")
|
||||
return True
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Download and extract arXiv LaTeX source files and PDFs only if both succeed.")
|
||||
parser.add_argument(
|
||||
"--url", type=str, default="https://arxiv.org/list/math/recent?skip=0&show=2000", help="URL of the arXiv list page to scrape (default: %(default)s)"
|
||||
)
|
||||
parser.add_argument("--data_dir", type=str, default="math_data/pdfs", help="Directory to save downloaded files (default: %(default)s)")
|
||||
parser.add_argument("--pdf_only", action="store_true", help="Skip LaTeX verification and only download PDFs")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.exists(args.data_dir):
|
||||
os.makedirs(args.data_dir)
|
||||
|
||||
print(f"Downloading list page from {args.url}...")
|
||||
response = requests.get(args.url)
|
||||
if response.status_code != 200:
|
||||
print(f"Error downloading list page: HTTP {response.status_code}")
|
||||
return
|
||||
|
||||
# Find all pdf links in the form: <a href="/pdf/2503.08675" ...>pdf</a>
|
||||
pattern = re.compile(r'href="/pdf/(\d+\.\d+)"')
|
||||
paper_ids = pattern.findall(response.text)
|
||||
print(f"Found {len(paper_ids)} papers.")
|
||||
|
||||
# For each paper, download based on the mode
|
||||
for paper_id in tqdm(paper_ids):
|
||||
if args.pdf_only:
|
||||
# Only download PDFs, skip LaTeX verification
|
||||
pdf_success = download_pdf(paper_id, args.data_dir)
|
||||
if not pdf_success:
|
||||
print(f"Failed to download PDF for {paper_id}")
|
||||
else:
|
||||
# Original behavior: only keep files if both tex extraction and pdf download succeed
|
||||
tex_success = download_and_extract_source(paper_id, args.data_dir)
|
||||
if not tex_success:
|
||||
print(f"Skipping PDF download for {paper_id} because tex extraction failed.")
|
||||
continue
|
||||
|
||||
pdf_success = download_pdf(paper_id, args.data_dir)
|
||||
if not pdf_success:
|
||||
# Remove the tex file if the PDF download fails.
|
||||
tex_path = os.path.join(args.data_dir, f"{paper_id}.tex")
|
||||
if os.path.exists(tex_path):
|
||||
os.remove(tex_path)
|
||||
print(f"Removed tex file for {paper_id} because PDF download failed.")
|
||||
time.sleep(3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,298 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_blank_pages_gpt.py - Identify PDF documents with blank pages and copy them.
|
||||
|
||||
This script:
|
||||
1. Takes a file containing S3 paths to PDF documents as input
|
||||
2. For each PDF, renders a random page and uses GPT-4o with the exact same query as buildsilver.py
|
||||
3. Identifies PDFs where the structured output has null natural_text
|
||||
4. Copies those PDF files to a new output folder
|
||||
|
||||
Usage:
|
||||
python mine_blank_pages_gpt.py --input_list path/to/s3_paths.txt --output_dir path/to/output --api_key your_openai_api_key
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Optional
|
||||
|
||||
import boto3
|
||||
import pypdf
|
||||
from openai import OpenAI
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
from olmocr.prompts import (
|
||||
build_openai_silver_data_prompt,
|
||||
openai_response_format_schema,
|
||||
)
|
||||
from olmocr.prompts.anchor import get_anchor_text
|
||||
|
||||
TARGET_IMAGE_DIM = 2048
|
||||
|
||||
|
||||
def download_pdf_from_s3(s3_path: str, local_path: str) -> bool:
|
||||
"""
|
||||
Download a PDF file from S3.
|
||||
|
||||
Args:
|
||||
s3_path: The S3 path (s3://bucket/path/to/file.pdf)
|
||||
local_path: The local path to save the file
|
||||
|
||||
Returns:
|
||||
bool: True if download was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Parse S3 path
|
||||
parts = s3_path.replace("s3://", "").split("/", 1)
|
||||
bucket = parts[0]
|
||||
key = parts[1]
|
||||
|
||||
# Create S3 client
|
||||
s3 = boto3.client("s3")
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
|
||||
# Download file
|
||||
s3.download_file(bucket, key, local_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {s3_path}: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def check_blank_page(pdf_path: str, page_num: int, api_key: str) -> Optional[bool]:
|
||||
"""
|
||||
Use GPT-4o with the exact same query as buildsilver.py to check if a page has null natural_text.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
Optional[bool]: True if natural_text is null, False otherwise, None if detection fails
|
||||
"""
|
||||
# Initialize OpenAI client
|
||||
client = OpenAI(api_key=api_key)
|
||||
|
||||
try:
|
||||
# Render the PDF page as an image (render_pdf_to_base64png is 1-indexed)
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=TARGET_IMAGE_DIM)
|
||||
|
||||
# Get anchor text
|
||||
anchor_text = get_anchor_text(pdf_path, page_num + 1, pdf_engine="pdfreport")
|
||||
|
||||
# Build the exact same prompt as buildsilver.py
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-4o-2024-08-06",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": build_openai_silver_data_prompt(anchor_text)},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
temperature=0.1,
|
||||
max_tokens=3000,
|
||||
logprobs=True,
|
||||
top_logprobs=5,
|
||||
response_format=openai_response_format_schema(),
|
||||
)
|
||||
|
||||
if not response.choices or len(response.choices) == 0:
|
||||
print(f"No response generated for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
# Parse the JSON response
|
||||
response_text = response.choices[0].message.content
|
||||
response_data = json.loads(response_text)
|
||||
|
||||
# Check if natural_text is null
|
||||
is_blank = response_data.get("natural_text") is None
|
||||
|
||||
if is_blank:
|
||||
print(f"Found blank page in {pdf_path} page {page_num + 1}")
|
||||
|
||||
return is_blank
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error checking {pdf_path} page {page_num}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def process_pdf(s3_path: str, temp_dir: str, output_dir: str, api_key: str) -> bool:
|
||||
"""
|
||||
Process a single PDF from S3.
|
||||
|
||||
Args:
|
||||
s3_path: S3 path to the PDF
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
bool: True if the PDF has a blank page and was copied, False otherwise
|
||||
"""
|
||||
# Extract filename from S3 path
|
||||
pdf_filename = os.path.basename(s3_path)
|
||||
local_pdf_path = os.path.join(temp_dir, pdf_filename)
|
||||
|
||||
# Download PDF from S3
|
||||
if not download_pdf_from_s3(s3_path, local_pdf_path):
|
||||
return False
|
||||
|
||||
pdf_filter = PdfFilter()
|
||||
|
||||
if pdf_filter.filter_out_pdf(local_pdf_path):
|
||||
print(f"Filtering out {pdf_filename}")
|
||||
return False
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(local_pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return False
|
||||
|
||||
# Select a random page to check
|
||||
page_num = random.randint(0, num_pages - 1)
|
||||
page_num = random.choice([page_num, 0]) # Bias 50% of the time to do the first page
|
||||
|
||||
# Check if the page has null natural_text
|
||||
is_blank = check_blank_page(local_pdf_path, page_num, api_key)
|
||||
|
||||
if is_blank:
|
||||
# Extract just the blank page and save it as a new PDF
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Create output filename with basename_pgnum.pdf format
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(output_dir, f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
|
||||
# Extract the single page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_pdf_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
print(f"Extracted blank page {page_num+1} from {pdf_filename} to {os.path.basename(output_pdf_path)}")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
return False
|
||||
finally:
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Identify and copy PDFs with blank pages")
|
||||
parser.add_argument("--input_list", required=True, help="Path to a file containing S3 paths to PDFs")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to copy PDFs with blank pages")
|
||||
parser.add_argument("--api_key", help="OpenAI API key (if not provided, will use OPENAI_API_KEY environment variable)")
|
||||
parser.add_argument("--temp_dir", default="/tmp/mine_blank_pages", help="Directory for temporary files")
|
||||
parser.add_argument("--max_pdfs", type=int, default=100, help="Maximum number of blank PDFs to find")
|
||||
parser.add_argument("--parallel", type=int, default=1, help="Number of parallel workers (default: 1 for sequential)")
|
||||
parser.add_argument("--reservoir_multiplier", type=int, default=100, help="Multiplier for reservoir sampling (default: 100x max_pdfs)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: OpenAI API key not provided. Use --api_key or set OPENAI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
os.makedirs(args.temp_dir, exist_ok=True)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Reservoir sampling to get random subset of PDFs
|
||||
reservoir_size = args.max_pdfs * args.reservoir_multiplier
|
||||
pdf_paths = []
|
||||
n = 0 # Total number of items seen
|
||||
|
||||
print(f"Using reservoir sampling with size {reservoir_size}")
|
||||
|
||||
with open(args.input_list, "r") as f:
|
||||
for line in f:
|
||||
n += 1
|
||||
path = line.strip()
|
||||
if not path:
|
||||
continue
|
||||
|
||||
if len(pdf_paths) < reservoir_size:
|
||||
pdf_paths.append(path)
|
||||
else:
|
||||
# Randomly decide whether to include this item
|
||||
s = random.randint(1, n)
|
||||
if s <= reservoir_size:
|
||||
pdf_paths[s - 1] = path
|
||||
|
||||
# Shuffle the reservoir
|
||||
random.shuffle(pdf_paths)
|
||||
|
||||
print(f"Sampled {len(pdf_paths)} PDF paths from {n} total paths")
|
||||
|
||||
blank_pdfs_found = 0
|
||||
|
||||
if args.parallel > 1:
|
||||
# Parallel processing
|
||||
print(f"Processing PDFs with {args.parallel} parallel workers")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=args.parallel) as executor:
|
||||
futures = []
|
||||
|
||||
# Submit all tasks
|
||||
for s3_path in pdf_paths:
|
||||
if blank_pdfs_found >= args.max_pdfs:
|
||||
break
|
||||
future = executor.submit(process_pdf, s3_path, args.temp_dir, args.output_dir, api_key)
|
||||
futures.append(future)
|
||||
|
||||
# Process results as they complete
|
||||
with tqdm(total=min(len(pdf_paths), args.max_pdfs), desc="Processing PDFs") as pbar:
|
||||
for future in as_completed(futures):
|
||||
try:
|
||||
result = future.result()
|
||||
if result:
|
||||
blank_pdfs_found += 1
|
||||
pbar.update(1)
|
||||
|
||||
if blank_pdfs_found >= args.max_pdfs:
|
||||
print(f"Reached maximum number of blank PDFs ({args.max_pdfs}), stopping")
|
||||
# Cancel remaining futures
|
||||
for f in futures:
|
||||
f.cancel()
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"Error in parallel processing: {str(e)}")
|
||||
else:
|
||||
# Sequential processing
|
||||
for s3_path in tqdm(pdf_paths, desc="Processing PDFs"):
|
||||
if process_pdf(s3_path, args.temp_dir, args.output_dir, api_key):
|
||||
blank_pdfs_found += 1
|
||||
|
||||
if blank_pdfs_found >= args.max_pdfs:
|
||||
print(f"Reached maximum number of blank PDFs ({args.max_pdfs}), stopping")
|
||||
break
|
||||
|
||||
print(f"Found and copied {blank_pdfs_found} PDFs with blank pages to {args.output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,226 @@
|
||||
import argparse
|
||||
import base64
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from collections import Counter
|
||||
from difflib import SequenceMatcher
|
||||
|
||||
import syntok.segmenter as segmenter
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
|
||||
from olmocr.bench.tests import TextPresenceTest, save_tests
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
|
||||
LABEL_WIDTH = 8 # fixed width for printing labels
|
||||
|
||||
# Uses a gemini prompt to get the most likely clean sentence from a pdf page
|
||||
last_gemini_call = time.perf_counter()
|
||||
|
||||
|
||||
def clean_base_sentence(pdf_path: str, page_num: int, base_sentence: str) -> str:
|
||||
client = genai.Client(
|
||||
api_key=os.environ.get("GEMINI_API_KEY"),
|
||||
)
|
||||
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=2048)
|
||||
image_part = types.Part(inline_data=types.Blob(mime_type="image/png", data=base64.b64decode(image_base64)))
|
||||
model = "gemini-2.0-flash-thinking-exp-01-21" # Consider using a more stable model for production
|
||||
# model="gemini-2.0-flash-001"
|
||||
contents = [
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
image_part,
|
||||
types.Part.from_text(
|
||||
text=f"""Base: {base_sentence}
|
||||
|
||||
Consider the sentence labeled "Base" above in the document image attached. What is the correct reading of this document within the image of the page? I need it to be exact down to the individual character and that's very important to get right. It needs to match the picture, not the provided text. Please just output the correct full sentence exactly how it appears in the document image and nothing else. You can merge hyphenated words back together, and don't output any new lines."""
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
generate_content_config = types.GenerateContentConfig(
|
||||
temperature=0.7,
|
||||
top_p=0.95,
|
||||
top_k=64,
|
||||
max_output_tokens=500,
|
||||
response_mime_type="text/plain",
|
||||
)
|
||||
|
||||
response = client.models.generate_content(
|
||||
model=model,
|
||||
contents=contents,
|
||||
config=generate_content_config,
|
||||
)
|
||||
|
||||
# Basic rate limitting
|
||||
global last_gemini_call
|
||||
if time.perf_counter() - last_gemini_call < 6:
|
||||
time.sleep(6 - (time.perf_counter() - last_gemini_call))
|
||||
|
||||
last_gemini_call = time.perf_counter()
|
||||
|
||||
# Return response
|
||||
if response is not None and response.candidates is not None and len(response.candidates) > 0:
|
||||
return response.candidates[0].content.parts[0].text
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def parse_sentences(text: str) -> list[str]:
|
||||
"""
|
||||
Splits a text into a list of sentence strings using syntok.
|
||||
Preserves original spacing and punctuation.
|
||||
"""
|
||||
sentences = []
|
||||
for paragraph in segmenter.process(text):
|
||||
for sentence in paragraph:
|
||||
# Reconstruct the sentence with original spacing
|
||||
sentence_str = ""
|
||||
for token in sentence:
|
||||
sentence_str += token.spacing + token.value
|
||||
# Trim any leading whitespace
|
||||
sentence_str = sentence_str.lstrip()
|
||||
sentences.append(sentence_str)
|
||||
return sentences
|
||||
|
||||
|
||||
def compare_votes_for_file(base_pdf_file: str, base_pdf_page: int, base_text: str, candidate_texts: list[str], max_diffs: int) -> None:
|
||||
"""
|
||||
For each sentence in the base text, finds the best matching sentence from
|
||||
each candidate text (using a similarity threshold). If any candidate sentences
|
||||
differ from the base sentence, collects that diff (base sentence plus variant
|
||||
votes) for later printing. At the end, prints only the top N diffs (by total vote count)
|
||||
for the file.
|
||||
|
||||
Comparison is case-insensitive, but output preserves original capitalization.
|
||||
"""
|
||||
base_sentences = parse_sentences(base_text)
|
||||
# Parse all candidate texts into lists of sentences
|
||||
candidate_sentences_list = [parse_sentences(ct) for ct in candidate_texts]
|
||||
|
||||
diffs = [] # list to hold diff entries
|
||||
for b_sentence in base_sentences:
|
||||
b_sentence = b_sentence.replace("\n", " ").strip()
|
||||
|
||||
votes = []
|
||||
for c_sentences in candidate_sentences_list:
|
||||
best_ratio = 0.0
|
||||
best_candidate = None
|
||||
|
||||
# Find the candidate sentence with the highest similarity to b_sentence
|
||||
# using case-insensitive comparison
|
||||
for c_sentence in c_sentences:
|
||||
ratio = SequenceMatcher(None, b_sentence.lower(), c_sentence.lower()).ratio()
|
||||
if ratio > best_ratio:
|
||||
best_ratio = ratio
|
||||
best_candidate = c_sentence # Keep original capitalization for output
|
||||
|
||||
# Append the candidate if it passes the similarity threshold (e.g., 0.7)
|
||||
if best_ratio > 0.5 and best_candidate is not None:
|
||||
votes.append(best_candidate.strip())
|
||||
|
||||
# Only consider variants that differ when compared case-insensitively
|
||||
variant_votes = [vote for vote in votes if vote.lower() != b_sentence.lower()]
|
||||
if variant_votes:
|
||||
diff_entry = {
|
||||
"base": b_sentence,
|
||||
"variants": Counter(variant_votes),
|
||||
"vote_count": len(variant_votes),
|
||||
}
|
||||
diffs.append(diff_entry)
|
||||
|
||||
# Sort diffs by vote_count descending and take only the top max_diffs
|
||||
diffs.sort(key=lambda d: d["vote_count"], reverse=True)
|
||||
top_diffs = diffs[:max_diffs]
|
||||
tests = []
|
||||
|
||||
for index, diff in enumerate(top_diffs):
|
||||
base_sentence = diff["base"]
|
||||
variant_counter = diff["variants"]
|
||||
|
||||
# Print base sentence using fixed-width label formatting
|
||||
print(f"{'Base:':<{LABEL_WIDTH}} {base_sentence}")
|
||||
print(f"{'Variants:':<{LABEL_WIDTH}}")
|
||||
for variant, count in variant_counter.items():
|
||||
label = f"{count}x:"
|
||||
print(f"{label:<{LABEL_WIDTH}} {variant}")
|
||||
# Get the clean version of the sentence
|
||||
cleaned = clean_base_sentence(base_pdf_file, base_pdf_page, base_sentence)
|
||||
print(f"{'Clean:':<{LABEL_WIDTH}} {cleaned}")
|
||||
print("-" * 40)
|
||||
|
||||
if cleaned is None:
|
||||
cleaned = base_sentence
|
||||
|
||||
tests.append(
|
||||
TextPresenceTest(
|
||||
pdf=os.path.basename(base_pdf_file),
|
||||
page=base_pdf_page,
|
||||
id=f"{os.path.basename(base_pdf_file).replace('.pdf', '')}_minediff_{index:02d}",
|
||||
type="present",
|
||||
threshold=1.0,
|
||||
text=cleaned,
|
||||
)
|
||||
)
|
||||
|
||||
return tests
|
||||
|
||||
|
||||
def get_pdf_from_md(md_path: str) -> str:
|
||||
base = os.path.basename(md_path)
|
||||
base = re.sub(r"_\d+\.md$", ".pdf", base)
|
||||
return os.path.join(os.path.dirname(md_path), "..", "pdfs", base)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Compares sentences from base and candidate texts, printing differences.")
|
||||
parser.add_argument("--base", default=os.path.join(os.path.dirname(__file__), "chatgpt"), help="Path to the folder containing base .md files.")
|
||||
parser.add_argument("--compare", default=os.path.join(os.path.dirname(__file__), "olmocr"), help="Path to the folder containing candidate .md files.")
|
||||
parser.add_argument("--max-diffs", type=int, default=5, help="Maximum number of diffs to display per file.")
|
||||
parser.add_argument(
|
||||
"--output", default="mine_diffs_candidates.jsonl", type=str, help="Output of potential candidate test proposals, to be verified or added to dataset"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
base_path = args.base
|
||||
compare_path = args.compare
|
||||
max_diffs = args.max_diffs
|
||||
|
||||
# Collect all .md files from the base and compare folders
|
||||
base_files = [f for f in os.listdir(base_path) if f.endswith(".md")]
|
||||
|
||||
all_tests = []
|
||||
|
||||
# Process each base file and print out the vote differences
|
||||
for bf in base_files:
|
||||
base_file_path = os.path.join(base_path, bf)
|
||||
with open(base_file_path, "r", encoding="utf-8") as f:
|
||||
base_text = f.read()
|
||||
|
||||
compare_files = [f for f in os.listdir(compare_path) if f.endswith(".md") and re.sub(r"_\d+\.md$", "", f) == re.sub(r"_\d+\.md$", "", bf)]
|
||||
|
||||
if not compare_files:
|
||||
print(f"skipping {bf} nothing to compare against")
|
||||
|
||||
# Read all candidate texts at once
|
||||
candidate_texts = []
|
||||
for cf in compare_files:
|
||||
with open(os.path.join(compare_path, cf), "r", encoding="utf-8") as f:
|
||||
candidate_texts.append(f.read())
|
||||
|
||||
base_pdf_file = get_pdf_from_md(base_file_path)
|
||||
base_pdf_page = 1
|
||||
print(f"Results for base file: {bf}")
|
||||
tests = compare_votes_for_file(base_pdf_file, base_pdf_page, base_text, candidate_texts, max_diffs)
|
||||
all_tests.extend(tests)
|
||||
print("")
|
||||
|
||||
# Output test candidates for review after each file, in case there are errors
|
||||
save_tests(all_tests, args.output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,305 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_footnotes_gpt.py - Identify PDF documents with footnotes and copy them.
|
||||
|
||||
This script:
|
||||
1. Takes a file containing S3 paths to PDF documents as input
|
||||
2. For each PDF, renders a random page and uses GPT-4o to check for footnotes
|
||||
3. Identifies PDFs where the page contains footnotes
|
||||
4. Copies those PDF files to a new output folder
|
||||
|
||||
Usage:
|
||||
python mine_footnotes_gpt.py --input_list path/to/s3_paths.txt --output_dir path/to/output --api_key your_openai_api_key
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import random
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import List, Optional
|
||||
|
||||
import boto3
|
||||
import pypdf
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
|
||||
TARGET_IMAGE_DIM = 1024
|
||||
|
||||
|
||||
class Footnote(BaseModel):
|
||||
marker: str
|
||||
text_before: Optional[str]
|
||||
text_after: Optional[str]
|
||||
|
||||
|
||||
class FootnoteDetectionResponse(BaseModel):
|
||||
"""Structured output for footnote detection."""
|
||||
|
||||
footnotes: List[Footnote]
|
||||
|
||||
|
||||
def download_pdf_from_s3(s3_path: str, local_path: str) -> bool:
|
||||
"""
|
||||
Download a PDF file from S3.
|
||||
|
||||
Args:
|
||||
s3_path: The S3 path (s3://bucket/path/to/file.pdf)
|
||||
local_path: The local path to save the file
|
||||
|
||||
Returns:
|
||||
bool: True if download was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Parse S3 path
|
||||
parts = s3_path.replace("s3://", "").split("/", 1)
|
||||
bucket = parts[0]
|
||||
key = parts[1]
|
||||
|
||||
# Create S3 client
|
||||
s3 = boto3.client("s3")
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
|
||||
# Download file
|
||||
s3.download_file(bucket, key, local_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {s3_path}: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def check_for_footnotes(pdf_path: str, page_num: int, api_key: str) -> Optional[bool]:
|
||||
"""
|
||||
Use GPT-4o to check if a page contains footnotes.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
Optional[bool]: True if page contains footnotes, False otherwise, None if detection fails
|
||||
"""
|
||||
# Initialize OpenAI client
|
||||
client = OpenAI(api_key=api_key)
|
||||
|
||||
try:
|
||||
# Render the PDF page as an image (render_pdf_to_base64png is 1-indexed)
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=TARGET_IMAGE_DIM)
|
||||
|
||||
# Simple prompt asking about footnotes
|
||||
prompt = (
|
||||
"Does this page contain footnotes? Ex. something you'd mark with a <sup> tag in html and is used to indicate some additional detail at the bottom of the page. "
|
||||
"Do not include references in this determination. The 'marker' is the little bit of text which is in a superscript. Output all the footnotes with their marker, and snippet of text occuring before or after the marker if present."
|
||||
)
|
||||
|
||||
response = client.beta.chat.completions.parse(
|
||||
model="gpt-5.1",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}}],
|
||||
}
|
||||
],
|
||||
temperature=1.0,
|
||||
max_completion_tokens=4000,
|
||||
response_format=FootnoteDetectionResponse,
|
||||
)
|
||||
|
||||
if not response.choices or len(response.choices) == 0:
|
||||
print(f"No response generated for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
# Parse the structured response
|
||||
parsed_response = response.choices[0].message.parsed
|
||||
|
||||
if parsed_response is None:
|
||||
print(f"Failed to parse response for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
# Check if there are any footnotes in the list
|
||||
footnotes = parsed_response.footnotes
|
||||
has_footnotes = len(footnotes) > 0
|
||||
|
||||
if has_footnotes:
|
||||
print(f"Found {len(footnotes)} footnote(s) in {pdf_path} page {page_num + 1}")
|
||||
|
||||
return has_footnotes
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error checking {pdf_path} page {page_num}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def process_pdf(s3_path: str, temp_dir: str, output_dir: str, api_key: str) -> bool:
|
||||
"""
|
||||
Process a single PDF from S3.
|
||||
|
||||
Args:
|
||||
s3_path: S3 path to the PDF
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
bool: True if the PDF has footnotes and was copied, False otherwise
|
||||
"""
|
||||
# Extract filename from S3 path
|
||||
pdf_filename = os.path.basename(s3_path)
|
||||
local_pdf_path = os.path.join(temp_dir, pdf_filename)
|
||||
|
||||
# Download PDF from S3
|
||||
if not download_pdf_from_s3(s3_path, local_pdf_path):
|
||||
return False
|
||||
|
||||
pdf_filter = PdfFilter()
|
||||
|
||||
if pdf_filter.filter_out_pdf(local_pdf_path):
|
||||
print(f"Filtering out {pdf_filename}")
|
||||
return False
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(local_pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return False
|
||||
|
||||
# Select a random page to check
|
||||
page_num = random.randint(0, num_pages - 1)
|
||||
page_num = random.choice([page_num, 0]) # Bias 50% of the time to do the first page
|
||||
|
||||
# Check if the page contains footnotes
|
||||
has_footnotes = check_for_footnotes(local_pdf_path, page_num, api_key)
|
||||
|
||||
if has_footnotes:
|
||||
# Extract just the page with footnotes and save it as a new PDF
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Create output filename with basename_pgnum.pdf format
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(output_dir, f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
|
||||
# Extract the single page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_pdf_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
print(f"Extracted page {page_num+1} with footnotes from {pdf_filename} to {os.path.basename(output_pdf_path)}")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
return False
|
||||
finally:
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Identify and copy PDFs with footnotes")
|
||||
parser.add_argument("--input_list", required=True, help="Path to a file containing S3 paths to PDFs")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to copy PDFs with footnotes")
|
||||
parser.add_argument("--api_key", help="OpenAI API key (if not provided, will use OPENAI_API_KEY environment variable)")
|
||||
parser.add_argument("--temp_dir", default="/tmp/mine_footnotes", help="Directory for temporary files")
|
||||
parser.add_argument("--max_pdfs", type=int, default=100, help="Maximum number of PDFs with footnotes to find")
|
||||
parser.add_argument("--parallel", type=int, default=1, help="Number of parallel workers (default: 1 for sequential)")
|
||||
parser.add_argument("--reservoir_multiplier", type=int, default=100, help="Multiplier for reservoir sampling (default: 100x max_pdfs)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: OpenAI API key not provided. Use --api_key or set OPENAI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
os.makedirs(args.temp_dir, exist_ok=True)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Reservoir sampling to get random subset of PDFs
|
||||
reservoir_size = args.max_pdfs * args.reservoir_multiplier
|
||||
pdf_paths = []
|
||||
n = 0 # Total number of items seen
|
||||
|
||||
print(f"Using reservoir sampling with size {reservoir_size}")
|
||||
|
||||
with open(args.input_list, "r") as f:
|
||||
for line in tqdm(f):
|
||||
n += 1
|
||||
path = line.strip()
|
||||
if not path:
|
||||
continue
|
||||
|
||||
if len(pdf_paths) < reservoir_size:
|
||||
pdf_paths.append(path)
|
||||
else:
|
||||
# Randomly decide whether to include this item
|
||||
s = random.randint(1, n)
|
||||
if s <= reservoir_size:
|
||||
pdf_paths[s - 1] = path
|
||||
|
||||
# Shuffle the reservoir
|
||||
random.shuffle(pdf_paths)
|
||||
|
||||
print(f"Sampled {len(pdf_paths)} PDF paths from {n} total paths")
|
||||
|
||||
footnote_pdfs_found = 0
|
||||
|
||||
if args.parallel > 1:
|
||||
# Parallel processing
|
||||
print(f"Processing PDFs with {args.parallel} parallel workers")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=args.parallel) as executor:
|
||||
futures = []
|
||||
|
||||
# Submit all tasks
|
||||
for s3_path in pdf_paths:
|
||||
if footnote_pdfs_found >= args.max_pdfs:
|
||||
break
|
||||
future = executor.submit(process_pdf, s3_path, args.temp_dir, args.output_dir, api_key)
|
||||
futures.append(future)
|
||||
|
||||
# Process results as they complete
|
||||
with tqdm(total=min(len(pdf_paths), args.max_pdfs), desc="Processing PDFs") as pbar:
|
||||
for future in as_completed(futures):
|
||||
try:
|
||||
result = future.result()
|
||||
if result:
|
||||
footnote_pdfs_found += 1
|
||||
pbar.update(1)
|
||||
|
||||
if footnote_pdfs_found >= args.max_pdfs:
|
||||
print(f"Reached maximum number of PDFs with footnotes ({args.max_pdfs}), stopping")
|
||||
# Cancel remaining futures
|
||||
for f in futures:
|
||||
f.cancel()
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"Error in parallel processing: {str(e)}")
|
||||
else:
|
||||
# Sequential processing
|
||||
for s3_path in tqdm(pdf_paths, desc="Processing PDFs"):
|
||||
if process_pdf(s3_path, args.temp_dir, args.output_dir, api_key):
|
||||
footnote_pdfs_found += 1
|
||||
|
||||
if footnote_pdfs_found >= args.max_pdfs:
|
||||
print(f"Reached maximum number of PDFs with footnotes ({args.max_pdfs}), stopping")
|
||||
break
|
||||
|
||||
print(f"Found and copied {footnote_pdfs_found} PDFs with footnotes to {args.output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,296 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_headers_footers.py - Extract headers and footers from PDF documents.
|
||||
|
||||
This script:
|
||||
1. Takes a file containing S3 paths to PDF documents as input
|
||||
2. For each PDF, extracts a random page and renders it to an image
|
||||
3. Uses Gemini to identify headers and footers in the rendered image
|
||||
4. Creates a test file asserting that the header/footer text should not appear
|
||||
5. Extracts the page from the PDF and saves it to an output folder
|
||||
|
||||
Usage:
|
||||
python mine_headers_footers.py --input_list path/to/s3_paths.txt --output_dir path/to/output --api_key your_gemini_api_key
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
from typing import List, Optional
|
||||
|
||||
import boto3
|
||||
import pypdf
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.bench.tests import TextPresenceTest, save_tests
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
|
||||
|
||||
def download_pdf_from_s3(s3_path: str, local_path: str) -> bool:
|
||||
"""
|
||||
Download a PDF file from S3.
|
||||
|
||||
Args:
|
||||
s3_path: The S3 path (s3://bucket/path/to/file.pdf)
|
||||
local_path: The local path to save the file
|
||||
|
||||
Returns:
|
||||
bool: True if download was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Parse S3 path
|
||||
parts = s3_path.replace("s3://", "").split("/", 1)
|
||||
bucket = parts[0]
|
||||
key = parts[1]
|
||||
|
||||
# Create S3 client
|
||||
s3 = boto3.client("s3")
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
|
||||
# Download file
|
||||
s3.download_file(bucket, key, local_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {s3_path}: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def extract_page_from_pdf(input_path: str, output_path: str, page_num: int) -> bool:
|
||||
"""
|
||||
Extract a specific page from a PDF and save it as a new PDF.
|
||||
|
||||
Args:
|
||||
input_path: Path to the input PDF
|
||||
output_path: Path to save the extracted page
|
||||
page_num: The page number to extract (0-indexed)
|
||||
|
||||
Returns:
|
||||
bool: True if extraction was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Ensure output directory exists
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
# Read the input PDF
|
||||
reader = pypdf.PdfReader(input_path)
|
||||
|
||||
# Check if page number is valid
|
||||
if page_num >= len(reader.pages):
|
||||
print(f"Page number {page_num} out of range for {input_path} with {len(reader.pages)} pages")
|
||||
return False
|
||||
|
||||
# Create a new PDF with just the selected page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error extracting page {page_num} from {input_path}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def detect_headers_footers(pdf_path: str, page_num: int, api_key: str) -> Optional[List[str]]:
|
||||
"""
|
||||
Use Gemini to detect headers and footers in a rendered PDF page.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: Gemini API key
|
||||
|
||||
Returns:
|
||||
Optional[List[str]]: List of detected header/footer texts, or None if detection failed
|
||||
"""
|
||||
client = genai.Client(
|
||||
api_key=os.environ.get("GEMINI_API_KEY"),
|
||||
)
|
||||
model = "gemini-2.0-flash"
|
||||
|
||||
# Render the PDF page as an image
|
||||
try:
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=2048) # render_pdf_to_base64png is 1-indexed
|
||||
except Exception as e:
|
||||
print(f"Error rendering PDF page: {str(e)}")
|
||||
return None
|
||||
|
||||
image_part = types.Part(inline_data=types.Blob(mime_type="image/png", data=base64.b64decode(image_base64)))
|
||||
|
||||
contents = [
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
image_part,
|
||||
types.Part.from_text(
|
||||
text="""Please extract and display the complete text from the document without omission. Include all sections and ensure nothing is summarized or abbreviated. I want the entire text of the document at any cost. Do not hallucinate"""
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
generate_content_config = types.GenerateContentConfig(
|
||||
temperature=1,
|
||||
top_p=0.95,
|
||||
top_k=40,
|
||||
max_output_tokens=8192,
|
||||
response_mime_type="application/json",
|
||||
response_schema=genai.types.Schema(
|
||||
type=genai.types.Type.OBJECT,
|
||||
properties={
|
||||
"headers": genai.types.Schema(
|
||||
type=genai.types.Type.ARRAY,
|
||||
items=genai.types.Schema(
|
||||
type=genai.types.Type.STRING,
|
||||
),
|
||||
),
|
||||
"footers": genai.types.Schema(
|
||||
type=genai.types.Type.ARRAY,
|
||||
items=genai.types.Schema(
|
||||
type=genai.types.Type.STRING,
|
||||
),
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
response = client.models.generate_content(model=model, contents=contents, config=generate_content_config)
|
||||
|
||||
assert len(response.candidates) > 0, "No candidates found"
|
||||
assert response.candidates[0].finish_reason == types.FinishReason.STOP, "Finish reason was not STOP, likely a processing error or repetition failure"
|
||||
|
||||
data = json.loads(response.candidates[0].content.parts[0].text)
|
||||
|
||||
return data.get("headers", []) + data.get("footers", [])
|
||||
|
||||
|
||||
def process_pdf(s3_path: str, temp_dir: str, output_dir: str, api_key: str, tests: List[TextPresenceTest]) -> None:
|
||||
"""
|
||||
Process a single PDF from S3.
|
||||
|
||||
Args:
|
||||
s3_path: S3 path to the PDF
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: Gemini API key
|
||||
tests: List to append tests to
|
||||
"""
|
||||
# Extract filename from S3 path
|
||||
pdf_filename = os.path.basename(s3_path)
|
||||
local_pdf_path = os.path.join(temp_dir, pdf_filename)
|
||||
|
||||
# Download PDF from S3
|
||||
if not download_pdf_from_s3(s3_path, local_pdf_path):
|
||||
return
|
||||
|
||||
pdf_filter = PdfFilter()
|
||||
|
||||
if pdf_filter.filter_out_pdf(local_pdf_path):
|
||||
print("Filtering out", pdf_filename)
|
||||
return
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(local_pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return
|
||||
|
||||
all_pages = list(range(len(reader.pages)))
|
||||
random.shuffle(all_pages)
|
||||
|
||||
for page_num in all_pages:
|
||||
# Detect headers and footers
|
||||
header_footer_text = detect_headers_footers(local_pdf_path, page_num, api_key)
|
||||
|
||||
# Only stick with headers and footers that have some actual data in them
|
||||
header_footer_text = [x for x in header_footer_text if len(x.strip()) > 3]
|
||||
|
||||
if not header_footer_text:
|
||||
print(f"No headers/footers detected in {pdf_filename} page {page_num}")
|
||||
continue
|
||||
|
||||
# Extract the page and save to output dir
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(output_dir, "pdfs", f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
|
||||
extract_page_from_pdf(local_pdf_path, output_pdf_path, page_num)
|
||||
|
||||
# TODO Now, process it again to make sure extracted headers/footers don't appear in the main body of the text
|
||||
|
||||
# Create tests for each header/footer text
|
||||
for i, text in enumerate(header_footer_text):
|
||||
test_id = f"{pdf_basename}_pg{page_num+1}_header_{i:02d}"
|
||||
test = TextPresenceTest(
|
||||
id=test_id,
|
||||
pdf=f"{pdf_basename}_pg{page_num+1}.pdf",
|
||||
page=1, # The extracted PDF has only one page
|
||||
type="absent",
|
||||
text=text,
|
||||
max_diffs=0,
|
||||
)
|
||||
tests.append(test)
|
||||
|
||||
print(f"Processed {pdf_filename} page {page_num+1}, found {len(header_footer_text)} headers/footers")
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
finally:
|
||||
# Clean up temporary file
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Extract headers and footers from PDF documents")
|
||||
parser.add_argument("--input_list", required=True, help="Path to a file containing S3 paths to PDFs")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to store extracted pages and tests")
|
||||
parser.add_argument("--api_key", help="Gemini API key (if not provided, will use GEMINI_API_KEY environment variable)")
|
||||
parser.add_argument("--temp_dir", default="/tmp/mine_headers_footers", help="Directory for temporary files")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("GEMINI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: Gemini API key not provided. Use --api_key or set GEMINI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
# Create directories
|
||||
os.makedirs(args.temp_dir, exist_ok=True)
|
||||
os.makedirs(os.path.join(args.output_dir, "pdfs"), exist_ok=True)
|
||||
|
||||
# Read input list
|
||||
with open(args.input_list, "r") as f:
|
||||
s3_paths = [line.strip() for line in f if line.strip()]
|
||||
|
||||
print(f"Found {len(s3_paths)} PDF paths in input list")
|
||||
|
||||
# Process each PDF
|
||||
tests = []
|
||||
for s3_path in tqdm(s3_paths, desc="Processing PDFs"):
|
||||
process_pdf(s3_path, args.temp_dir, args.output_dir, api_key, tests)
|
||||
|
||||
# Save tests after each PDF to avoid losing data in case of crashes
|
||||
if tests:
|
||||
save_tests(tests, os.path.join(args.output_dir, "header_footer_tests.jsonl"))
|
||||
|
||||
if len(tests) > 100:
|
||||
break
|
||||
|
||||
print(f"Saved {len(tests)} tests to {os.path.join(args.output_dir, 'header_footer_tests.jsonl')}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,341 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_length_gpt_simple.py - Identify PDF documents by word count and copy them.
|
||||
|
||||
This script:
|
||||
1. Takes a file containing S3 paths to PDF documents as input
|
||||
2. For each PDF, renders a random page and uses GPT to analyze document structure
|
||||
3. Identifies document elements and estimates word counts
|
||||
4. Copies those PDF pages to folders organized by total word count
|
||||
|
||||
Usage:
|
||||
python mine_length_gpt_simple.py --input_list path/to/s3_paths.txt --output_dir path/to/output --api_key your_openai_api_key
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import random
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Optional
|
||||
|
||||
import boto3
|
||||
import pypdf
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
|
||||
TARGET_IMAGE_DIM = 1024
|
||||
|
||||
|
||||
class StructureElement(BaseModel):
|
||||
"""Information about a single document element."""
|
||||
|
||||
type: str # One of: paragraph, heading, footer, table, equation, image
|
||||
estimated_words: int
|
||||
|
||||
|
||||
class DocumentStructure(BaseModel):
|
||||
"""Structured output for document structure analysis."""
|
||||
|
||||
elements: list[StructureElement]
|
||||
|
||||
|
||||
def download_pdf_from_s3(s3_path: str, local_path: str) -> bool:
|
||||
"""
|
||||
Download a PDF file from S3.
|
||||
|
||||
Args:
|
||||
s3_path: The S3 path (s3://bucket/path/to/file.pdf)
|
||||
local_path: The local path to save the file
|
||||
|
||||
Returns:
|
||||
bool: True if download was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Parse S3 path
|
||||
parts = s3_path.replace("s3://", "").split("/", 1)
|
||||
bucket = parts[0]
|
||||
key = parts[1]
|
||||
|
||||
# Create S3 client
|
||||
s3 = boto3.client("s3")
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
|
||||
# Download file
|
||||
s3.download_file(bucket, key, local_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {s3_path}: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def get_word_count_bucket(total_words: int) -> str:
|
||||
"""
|
||||
Get the folder name for a given word count, bucketed by powers of 2.
|
||||
|
||||
Args:
|
||||
total_words: Total number of words across all elements
|
||||
|
||||
Returns:
|
||||
str: Folder name like "0_words", "1_word", "2_words", "4_words", etc.
|
||||
"""
|
||||
if total_words == 0:
|
||||
return "0_words"
|
||||
elif total_words == 1:
|
||||
return "1_word"
|
||||
else:
|
||||
# Find the next power of 2 >= total_words
|
||||
power = 1
|
||||
while power < total_words:
|
||||
power *= 2
|
||||
return f"{power}_words"
|
||||
|
||||
|
||||
def analyze_document_structure(pdf_path: str, page_num: int, api_key: str) -> Optional[int]:
|
||||
"""
|
||||
Use GPT to analyze document structure and estimate word count.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
Optional[int]: Total estimated word count, or None if analysis fails
|
||||
"""
|
||||
# Initialize OpenAI client
|
||||
client = OpenAI(api_key=api_key)
|
||||
|
||||
try:
|
||||
# Render the PDF page as an image (render_pdf_to_base64png is 1-indexed)
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=TARGET_IMAGE_DIM)
|
||||
|
||||
# Prompt asking for detailed document structure
|
||||
prompt = """Analyze the structure of this document page. Identify all distinct elements on the page and classify each one.
|
||||
|
||||
For each element, specify:
|
||||
- type: One of the following: paragraph, heading, footer, table, equation, image
|
||||
- estimated_words: Your best estimate of the number of words of text in that element
|
||||
|
||||
Be thorough and identify all visible elements."""
|
||||
|
||||
response = client.beta.chat.completions.parse(
|
||||
model="gpt-5.1",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}}],
|
||||
}
|
||||
],
|
||||
max_completion_tokens=2000,
|
||||
response_format=DocumentStructure,
|
||||
)
|
||||
|
||||
if not response.choices or len(response.choices) == 0:
|
||||
print(f"No response generated for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
# Parse the structured response
|
||||
parsed_response = response.choices[0].message.parsed
|
||||
|
||||
if parsed_response is None:
|
||||
print(f"Failed to parse response for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
elements = parsed_response.elements
|
||||
total_words = sum(element.estimated_words for element in elements)
|
||||
|
||||
print(f"Found {len(elements)} element(s) in {pdf_path} page {page_num + 1}, total words: {total_words}")
|
||||
|
||||
# Group elements by type and show summary
|
||||
element_summary = {}
|
||||
for element in elements:
|
||||
if element.type not in element_summary:
|
||||
element_summary[element.type] = {"count": 0, "words": 0}
|
||||
element_summary[element.type]["count"] += 1
|
||||
element_summary[element.type]["words"] += element.estimated_words
|
||||
|
||||
for element_type, stats in sorted(element_summary.items()):
|
||||
print(f" {element_type}: {stats['count']} element(s), {stats['words']} words")
|
||||
|
||||
return total_words
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error analyzing {pdf_path} page {page_num}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def process_pdf(s3_path: str, temp_dir: str, output_dir: str, api_key: str) -> bool:
|
||||
"""
|
||||
Process a single PDF from S3.
|
||||
|
||||
Args:
|
||||
s3_path: S3 path to the PDF
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
bool: True if the PDF was analyzed and copied, False otherwise
|
||||
"""
|
||||
# Extract filename from S3 path
|
||||
pdf_filename = os.path.basename(s3_path)
|
||||
local_pdf_path = os.path.join(temp_dir, pdf_filename)
|
||||
|
||||
# Download PDF from S3
|
||||
if not download_pdf_from_s3(s3_path, local_pdf_path):
|
||||
return False
|
||||
|
||||
pdf_filter = PdfFilter()
|
||||
|
||||
if pdf_filter.filter_out_pdf(local_pdf_path):
|
||||
print(f"Filtering out {pdf_filename}")
|
||||
return False
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(local_pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return False
|
||||
|
||||
# Select a random page to check
|
||||
page_num = random.randint(0, num_pages - 1)
|
||||
page_num = random.choice([page_num, 0]) # Bias 50% of the time to do the first page
|
||||
|
||||
# Analyze document structure
|
||||
total_words = analyze_document_structure(local_pdf_path, page_num, api_key)
|
||||
|
||||
if total_words is None:
|
||||
return False
|
||||
|
||||
# Get the word count bucket for organizing output
|
||||
bucket_name = get_word_count_bucket(total_words)
|
||||
bucket_dir = os.path.join(output_dir, bucket_name)
|
||||
os.makedirs(bucket_dir, exist_ok=True)
|
||||
|
||||
# Create output filename with basename_pgnum.pdf format
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(bucket_dir, f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
|
||||
# Extract the single page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_pdf_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
print(f"Extracted page {page_num+1} from {pdf_filename} to {bucket_name}/{os.path.basename(output_pdf_path)}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
return False
|
||||
finally:
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Identify and copy PDFs by word count")
|
||||
parser.add_argument("--input_list", required=True, help="Path to a file containing S3 paths to PDFs")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to copy PDF pages")
|
||||
parser.add_argument("--api_key", help="OpenAI API key (if not provided, will use OPENAI_API_KEY environment variable)")
|
||||
parser.add_argument("--temp_dir", default="/tmp/mine_length", help="Directory for temporary files")
|
||||
parser.add_argument("--max_pdfs", type=int, default=100, help="Maximum number of PDFs to process")
|
||||
parser.add_argument("--parallel", type=int, default=1, help="Number of parallel workers (default: 1 for sequential)")
|
||||
parser.add_argument("--reservoir_multiplier", type=int, default=100, help="Multiplier for reservoir sampling (default: 100x max_pdfs)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: OpenAI API key not provided. Use --api_key or set OPENAI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
os.makedirs(args.temp_dir, exist_ok=True)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Reservoir sampling to get random subset of PDFs
|
||||
reservoir_size = args.max_pdfs * args.reservoir_multiplier
|
||||
pdf_paths = []
|
||||
n = 0 # Total number of items seen
|
||||
|
||||
print(f"Using reservoir sampling with size {reservoir_size}")
|
||||
|
||||
with open(args.input_list, "r") as f:
|
||||
for line in tqdm(f):
|
||||
n += 1
|
||||
path = line.strip()
|
||||
if not path:
|
||||
continue
|
||||
|
||||
if len(pdf_paths) < reservoir_size:
|
||||
pdf_paths.append(path)
|
||||
else:
|
||||
# Randomly decide whether to include this item
|
||||
s = random.randint(1, n)
|
||||
if s <= reservoir_size:
|
||||
pdf_paths[s - 1] = path
|
||||
|
||||
# Shuffle the reservoir
|
||||
random.shuffle(pdf_paths)
|
||||
|
||||
print(f"Sampled {len(pdf_paths)} PDF paths from {n} total paths")
|
||||
|
||||
pdfs_processed = 0
|
||||
|
||||
if args.parallel > 1:
|
||||
# Parallel processing
|
||||
print(f"Processing PDFs with {args.parallel} parallel workers")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=args.parallel) as executor:
|
||||
futures = []
|
||||
|
||||
# Submit all tasks
|
||||
for s3_path in pdf_paths:
|
||||
if pdfs_processed >= args.max_pdfs:
|
||||
break
|
||||
future = executor.submit(process_pdf, s3_path, args.temp_dir, args.output_dir, api_key)
|
||||
futures.append(future)
|
||||
|
||||
# Process results as they complete
|
||||
with tqdm(total=min(len(pdf_paths), args.max_pdfs), desc="Processing PDFs") as pbar:
|
||||
for future in as_completed(futures):
|
||||
try:
|
||||
result = future.result()
|
||||
if result:
|
||||
pdfs_processed += 1
|
||||
pbar.update(1)
|
||||
|
||||
if pdfs_processed >= args.max_pdfs:
|
||||
print(f"Reached maximum number of PDFs ({args.max_pdfs}), stopping")
|
||||
# Cancel remaining futures
|
||||
for f in futures:
|
||||
f.cancel()
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"Error in parallel processing: {str(e)}")
|
||||
else:
|
||||
# Sequential processing
|
||||
for s3_path in tqdm(pdf_paths, desc="Processing PDFs"):
|
||||
if process_pdf(s3_path, args.temp_dir, args.output_dir, api_key):
|
||||
pdfs_processed += 1
|
||||
|
||||
if pdfs_processed >= args.max_pdfs:
|
||||
print(f"Reached maximum number of PDFs ({args.max_pdfs}), stopping")
|
||||
break
|
||||
|
||||
print(f"Processed and copied {pdfs_processed} PDF pages to {args.output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,314 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_long_tiny_text.py - Extract long text from PDF documents.
|
||||
|
||||
This script:
|
||||
1. Takes a folder containing PDF documents as input
|
||||
2. For each PDF, extracts random pages and renders them to images
|
||||
3. Uses Gemini to identify text in the rendered images
|
||||
4. Creates test files asserting that the text should be present
|
||||
5. Extracts the pages from the PDF and saves them to an output folder
|
||||
|
||||
Usage:
|
||||
python mine_long_tiny_text.py --input_dir path/to/pdf_folder --output_dir path/to/output --api_key your_gemini_api_key
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
from typing import List
|
||||
|
||||
import pypdf
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.bench.tests import TextPresenceTest, save_tests
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
|
||||
|
||||
def extract_page_from_pdf(input_path: str, output_path: str, page_num: int) -> bool:
|
||||
"""
|
||||
Extract a specific page from a PDF and save it as a new PDF.
|
||||
|
||||
Args:
|
||||
input_path: Path to the input PDF
|
||||
output_path: Path to save the extracted page
|
||||
page_num: The page number to extract (0-indexed)
|
||||
|
||||
Returns:
|
||||
bool: True if extraction was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Ensure output directory exists
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
# Read the input PDF
|
||||
reader = pypdf.PdfReader(input_path)
|
||||
|
||||
# Check if page number is valid
|
||||
if page_num >= len(reader.pages):
|
||||
print(f"Page number {page_num} out of range for {input_path} with {len(reader.pages)} pages")
|
||||
return False
|
||||
|
||||
# Create a new PDF with just the selected page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error extracting page {page_num} from {input_path}: {str(e)}")
|
||||
return False # Return False instead of raising to continue processing other PDFs
|
||||
|
||||
|
||||
def detect_long_text(pdf_path: str, page_num: int, api_key: str) -> List[str]:
|
||||
"""
|
||||
Use Gemini to detect long text in a rendered PDF page.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: Gemini API key
|
||||
|
||||
Returns:
|
||||
List[str]: List of detected text, empty list if detection failed
|
||||
"""
|
||||
try:
|
||||
client = genai.Client(
|
||||
api_key=api_key, # Use the provided API key instead of environment variable
|
||||
)
|
||||
model = "gemini-2.0-flash"
|
||||
|
||||
# Render the PDF page as an image
|
||||
try:
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=2048) # render_pdf_to_base64png is 1-indexed
|
||||
except Exception as e:
|
||||
print(f"Error rendering PDF page {page_num+1} from {pdf_path}: {str(e)}")
|
||||
return []
|
||||
|
||||
image_part = types.Part(inline_data=types.Blob(mime_type="image/png", data=base64.b64decode(image_base64)))
|
||||
|
||||
contents = [
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
image_part,
|
||||
types.Part.from_text(
|
||||
text="""Extract and display all text from the document without any omissions. The documents may be in a multi-column format, so handle that carefully. If the words are less than 9, combine text from the next line. I don't want all the text, just randomly pick few sentences. Do not summarize, abbreviate, or hallucinate any content."""
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
generate_content_config = types.GenerateContentConfig(
|
||||
temperature=1,
|
||||
top_p=0.95,
|
||||
top_k=40,
|
||||
max_output_tokens=8192,
|
||||
response_mime_type="application/json",
|
||||
response_schema=genai.types.Schema(
|
||||
type=genai.types.Type.OBJECT,
|
||||
properties={
|
||||
"equations": genai.types.Schema(
|
||||
type=genai.types.Type.ARRAY,
|
||||
items=genai.types.Schema(
|
||||
type=genai.types.Type.STRING,
|
||||
),
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
response = client.models.generate_content(model=model, contents=contents, config=generate_content_config)
|
||||
|
||||
if len(response.candidates) == 0:
|
||||
print(f"No candidates found for {pdf_path} page {page_num+1}")
|
||||
return []
|
||||
|
||||
if response.candidates[0].finish_reason != types.FinishReason.STOP:
|
||||
print(f"Finish reason was not STOP for {pdf_path} page {page_num+1}, likely a processing error")
|
||||
return []
|
||||
|
||||
try:
|
||||
data = json.loads(response.candidates[0].content.parts[0].text)
|
||||
return data.get("equations", [])
|
||||
except json.JSONDecodeError:
|
||||
print(f"Failed to parse JSON response for {pdf_path} page {page_num+1}")
|
||||
return []
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error detecting text in {pdf_path} page {page_num+1}: {str(e)}")
|
||||
return []
|
||||
|
||||
|
||||
def process_pdf(
|
||||
pdf_path: str, output_dir: str, api_key: str, tests: List[TextPresenceTest], max_pages_per_pdf: int = 20, force_processing: bool = True
|
||||
) -> bool:
|
||||
"""
|
||||
Process a single PDF, extracting text from multiple pages.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF
|
||||
output_dir: Directory for output files
|
||||
api_key: Gemini API key
|
||||
tests: List to append tests to
|
||||
max_pages_per_pdf: Maximum number of pages to process per PDF
|
||||
force_processing: If True, process PDF even if it would normally be filtered out
|
||||
|
||||
Returns:
|
||||
bool: True if PDF was processed successfully
|
||||
"""
|
||||
# Extract filename from path
|
||||
pdf_filename = os.path.basename(pdf_path)
|
||||
|
||||
pdf_filter = PdfFilter()
|
||||
|
||||
if not force_processing and pdf_filter.filter_out_pdf(pdf_path):
|
||||
print(f"Filtering out {pdf_filename} (use --force_processing to override)")
|
||||
return False
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return False
|
||||
|
||||
# Get all pages and shuffle them to select a random subset
|
||||
all_pages = list(range(num_pages))
|
||||
random.shuffle(all_pages)
|
||||
|
||||
# Take only the specified maximum number of pages
|
||||
pages_to_process = all_pages[: min(max_pages_per_pdf, num_pages)]
|
||||
|
||||
processed_pages = 0
|
||||
pdf_processed = False # Flag to track if at least one page was processed
|
||||
|
||||
for page_num in pages_to_process:
|
||||
# Detect text
|
||||
text_sections = detect_long_text(pdf_path, page_num, api_key)
|
||||
|
||||
# Only keep text sections that are non-empty
|
||||
text_sections = [text for text in text_sections if len(text.strip()) > 3]
|
||||
|
||||
# Extract the page regardless of text detection
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(output_dir, "pdfs", f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
|
||||
# Extract the page regardless of text detection
|
||||
page_extracted = extract_page_from_pdf(pdf_path, output_pdf_path, page_num)
|
||||
if page_extracted:
|
||||
pdf_processed = True
|
||||
|
||||
if not text_sections:
|
||||
print(f"No text detected in {pdf_filename} page {page_num+1} - creating empty test")
|
||||
# Create a placeholder test to ensure the PDF is represented
|
||||
test_id = f"{pdf_basename}_pg{page_num+1}_no_text"
|
||||
test = TextPresenceTest(
|
||||
id=test_id,
|
||||
pdf=f"{pdf_basename}_pg{page_num+1}.pdf",
|
||||
page=1, # The extracted PDF has only one page
|
||||
type="present",
|
||||
text="No text detected", # Placeholder text
|
||||
max_diffs=0,
|
||||
)
|
||||
tests.append(test)
|
||||
else:
|
||||
# Create tests for each text section
|
||||
for i, text_section in enumerate(text_sections):
|
||||
test_id = f"{pdf_basename}_pg{page_num+1}_text_{i:02d}"
|
||||
test = TextPresenceTest(
|
||||
id=test_id,
|
||||
pdf=f"{pdf_basename}_pg{page_num+1}.pdf",
|
||||
page=1, # The extracted PDF has only one page
|
||||
type="present",
|
||||
text=text_section,
|
||||
max_diffs=0,
|
||||
)
|
||||
tests.append(test)
|
||||
|
||||
print(f"Processed {pdf_filename} page {page_num+1}, found {len(text_sections)} text sections")
|
||||
processed_pages += 1
|
||||
|
||||
print(f"Completed processing {processed_pages} pages from {pdf_filename}")
|
||||
return pdf_processed
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def get_pdf_files_from_directory(directory: str) -> List[str]:
|
||||
"""
|
||||
Get a list of all PDF files in a directory.
|
||||
|
||||
Args:
|
||||
directory: Path to the directory containing PDFs
|
||||
|
||||
Returns:
|
||||
List[str]: List of full paths to PDF files
|
||||
"""
|
||||
pdf_files = []
|
||||
|
||||
for filename in os.listdir(directory):
|
||||
if filename.lower().endswith(".pdf"):
|
||||
full_path = os.path.join(directory, filename)
|
||||
if os.path.isfile(full_path):
|
||||
pdf_files.append(full_path)
|
||||
|
||||
return pdf_files
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Extract long text from PDF documents")
|
||||
parser.add_argument("--input_dir", required=True, help="Directory containing PDF files")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to store extracted pages and tests")
|
||||
parser.add_argument("--api_key", help="Gemini API key (if not provided, will use GEMINI_API_KEY environment variable)")
|
||||
parser.add_argument("--force_processing", action="store_true", help="Process all PDFs even if they would normally be filtered out")
|
||||
parser.add_argument("--max_pages_per_pdf", type=int, default=20, help="Maximum number of pages to process per PDF")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("GEMINI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: Gemini API key not provided. Use --api_key or set GEMINI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
# Create directories
|
||||
os.makedirs(os.path.join(args.output_dir, "pdfs"), exist_ok=True)
|
||||
|
||||
# Get PDF files from input directory
|
||||
pdf_files = get_pdf_files_from_directory(args.input_dir)
|
||||
|
||||
if not pdf_files:
|
||||
print(f"No PDF files found in {args.input_dir}")
|
||||
return
|
||||
|
||||
print(f"Found {len(pdf_files)} PDF files in input directory")
|
||||
|
||||
# Process each PDF
|
||||
tests = []
|
||||
processed_pdfs = 0
|
||||
for pdf_path in tqdm(pdf_files, desc="Processing PDFs"):
|
||||
if process_pdf(pdf_path, args.output_dir, api_key, tests, args.max_pages_per_pdf, args.force_processing):
|
||||
processed_pdfs += 1
|
||||
|
||||
# Save tests after each PDF to avoid losing data in case of crashes
|
||||
if tests:
|
||||
save_tests(tests, os.path.join(args.output_dir, "long_tests.jsonl"))
|
||||
|
||||
print(f"Successfully processed {processed_pdfs} out of {len(pdf_files)} PDFs")
|
||||
print(f"Saved {len(tests)} tests to {os.path.join(args.output_dir, 'long_tests.jsonl')}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,373 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_math.py - Extract and validate math equations from candidate files and TeX bases.
|
||||
|
||||
This upgraded version:
|
||||
• Uses the Python logging module for cleaner logging.
|
||||
• Uses tqdm to display a progress bar.
|
||||
• Uses ProcessPoolExecutor to process TeX file groups in parallel.
|
||||
• For each TeX file, shuffles its pages randomly and processes them one-by-one.
|
||||
Once three pages return at least one equation each, further pages are skipped.
|
||||
• Adds an argparse argument for the similarity threshold for matches.
|
||||
• Saves JSONL outputs incrementally as each TeX file group is processed.
|
||||
|
||||
Usage:
|
||||
python mine_math.py --math_data /path/to/math_data --candidate candidate_folder --output_file math_tests.jsonl
|
||||
[--max_pages 3] [--parallel 8] [--sim_threshold 0.7]
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import numba
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.bench.katex.render import render_equation
|
||||
from olmocr.bench.tests import (
|
||||
MathTest, # Assumes MathTest is JSON serializable or has __dict__
|
||||
)
|
||||
from olmocr.bench.tests import (
|
||||
save_tests, # Original saving function (not used for incremental save)
|
||||
)
|
||||
|
||||
# --- Logging Setup ---
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
|
||||
|
||||
# --- Utility Functions ---
|
||||
|
||||
|
||||
def normalize_text(text: str) -> str:
|
||||
"""Normalize text for better matching."""
|
||||
text = re.sub(r"\s+", " ", text)
|
||||
replacements = {"'": "'", "‚": "'", '"': '"', "„": '"', "＿": "_", "–": "-", "—": "-", "‑": "-", "‒": "-"}
|
||||
for fancy_char, ascii_char in replacements.items():
|
||||
text = text.replace(fancy_char, ascii_char)
|
||||
return text
|
||||
|
||||
|
||||
def extract_tex_content(tex_file: str) -> str:
|
||||
"""Extract the content from a TeX file."""
|
||||
try:
|
||||
with open(tex_file, "r", encoding="utf-8") as f:
|
||||
return f.read()
|
||||
except UnicodeDecodeError:
|
||||
try:
|
||||
with open(tex_file, "r", encoding="latin-1") as f:
|
||||
return f.read()
|
||||
except Exception as e:
|
||||
logging.error("Error reading %s: %s", tex_file, e)
|
||||
return ""
|
||||
|
||||
|
||||
def extract_candidate_content(candidate_file: str) -> str:
|
||||
"""Extract the content from a candidate .md file."""
|
||||
try:
|
||||
with open(candidate_file, "r", encoding="utf-8") as f:
|
||||
return f.read()
|
||||
except Exception as e:
|
||||
logging.error("Error reading %s: %s", candidate_file, e)
|
||||
return ""
|
||||
|
||||
|
||||
def extract_math_from_tex(tex_content: str) -> List[Tuple[str, str]]:
|
||||
"""
|
||||
Extract math equations from TeX content.
|
||||
Returns list of tuples (equation_type, equation_content)
|
||||
"""
|
||||
math_equations = []
|
||||
|
||||
# Patterns for display math
|
||||
display_patterns = [
|
||||
(r"\$\$(.*?)\$\$", "$$"),
|
||||
(r"\\begin\{equation\}(.*?)\\end\{equation\}", "equation"),
|
||||
(r"\\begin\{equation\*\}(.*?)\\end\{equation\*\}", "equation*"),
|
||||
(r"\\begin\{align\}(.*?)\\end\{align\}", "align"),
|
||||
(r"\\begin\{align\*\}(.*?)\\end\{align\*\}", "align*"),
|
||||
(r"\\begin\{displaymath\}(.*?)\\end\{displaymath\}", "displaymath"),
|
||||
(r"\\\[(.*?)\\\]", "displaymath"),
|
||||
]
|
||||
# Patterns for inline math
|
||||
inline_patterns = [(r"\$(.*?)\$", "inline"), (r"\\\((.*?)\\\)", "inline")]
|
||||
|
||||
for pattern_list in [display_patterns, inline_patterns]:
|
||||
for pattern, eq_type in pattern_list:
|
||||
matches = re.finditer(pattern, tex_content, re.DOTALL)
|
||||
for match in matches:
|
||||
equation = match.group(1).strip()
|
||||
if equation and not equation.isspace():
|
||||
math_equations.append((eq_type, equation))
|
||||
return math_equations
|
||||
|
||||
|
||||
@numba.njit
|
||||
def compute_dp(candidate_arr, text_arr):
|
||||
m = candidate_arr.shape[0]
|
||||
n = text_arr.shape[0]
|
||||
dp = np.empty((m + 1, n + 1), dtype=np.int32)
|
||||
# For empty candidate, cost is 0 (can match anywhere in text)
|
||||
for j in range(n + 1):
|
||||
dp[0, j] = 0
|
||||
# When text is empty, need to delete all candidate characters.
|
||||
for i in range(1, m + 1):
|
||||
dp[i, 0] = i
|
||||
|
||||
for i in range(1, m + 1):
|
||||
for j in range(1, n + 1):
|
||||
cost = 0 if candidate_arr[i - 1] == text_arr[j - 1] else 1
|
||||
dp[i, j] = min(
|
||||
dp[i - 1, j - 1] + cost, dp[i - 1, j] + 1, dp[i, j - 1] + 1 # substitution or match # deletion (from candidate)
|
||||
) # insertion (in candidate)
|
||||
return dp
|
||||
|
||||
|
||||
@numba.njit
|
||||
def find_best_end(dp, m, n):
|
||||
best_distance = 1 << 30 # a large number
|
||||
best_end = 0
|
||||
for j in range(n + 1):
|
||||
if dp[m, j] < best_distance:
|
||||
best_distance = dp[m, j]
|
||||
best_end = j
|
||||
return best_end, best_distance
|
||||
|
||||
|
||||
@numba.njit
|
||||
def backtrack(dp, candidate_arr, text_arr, m, best_end):
|
||||
i = m
|
||||
j = best_end
|
||||
while i > 0:
|
||||
# Check for a diagonal move (match or substitution)
|
||||
if j > 0 and dp[i, j] == dp[i - 1, j - 1] + (0 if candidate_arr[i - 1] == text_arr[j - 1] else 1):
|
||||
i -= 1
|
||||
j -= 1
|
||||
elif dp[i, j] == dp[i - 1, j] + 1:
|
||||
i -= 1
|
||||
else:
|
||||
j -= 1
|
||||
return j # start index in text
|
||||
|
||||
|
||||
def find_matching_content(candidate_text: str, tex_content: str, sim_threshold: float) -> Optional[str]:
|
||||
"""
|
||||
Find the substring of tex_content that most closely matches candidate_text using
|
||||
dynamic programming accelerated by numba. Returns the matching substring if its
|
||||
normalized similarity (1 - (edit_distance / len(candidate_text))) is above sim_threshold,
|
||||
otherwise returns None.
|
||||
"""
|
||||
candidate_norm = normalize_text(candidate_text)
|
||||
tex_norm = normalize_text(tex_content)
|
||||
|
||||
m = len(candidate_norm)
|
||||
n = len(tex_norm)
|
||||
if m == 0 or n == 0:
|
||||
return None
|
||||
|
||||
# Convert strings to numpy arrays of integer character codes.
|
||||
candidate_arr = np.empty(m, dtype=np.int32)
|
||||
for i, c in enumerate(candidate_norm):
|
||||
candidate_arr[i] = ord(c)
|
||||
text_arr = np.empty(n, dtype=np.int32)
|
||||
for j, c in enumerate(tex_norm):
|
||||
text_arr[j] = ord(c)
|
||||
|
||||
dp = compute_dp(candidate_arr, text_arr)
|
||||
best_end, min_distance = find_best_end(dp, m, n)
|
||||
similarity = (m - min_distance) / m
|
||||
|
||||
logging.info("Similarity: %.3f", similarity)
|
||||
if similarity < sim_threshold:
|
||||
return None
|
||||
start_index = backtrack(dp, candidate_arr, text_arr, m, best_end)
|
||||
return tex_norm[start_index:best_end]
|
||||
|
||||
|
||||
def parse_candidate_filename(filename: str) -> Optional[Tuple[str, int]]:
|
||||
"""
|
||||
Parse candidate filename in the format: [tex file basename]_pg[pagenum]_repeat1.md
|
||||
Returns tuple (tex_basename, page_num) or None if the format doesn't match.
|
||||
"""
|
||||
basename = os.path.basename(filename)
|
||||
match = re.match(r"(.+)_pg(\d+)_repeat\d+\.md$", basename)
|
||||
if match:
|
||||
tex_basename = match.group(1)
|
||||
page_num = int(match.group(2))
|
||||
return tex_basename, page_num
|
||||
return None
|
||||
|
||||
|
||||
def validate_equation(equation: str) -> bool:
|
||||
"""
|
||||
Validate that an equation renders correctly with KaTeX.
|
||||
Returns True if the equation is valid, False otherwise.
|
||||
"""
|
||||
rendered = render_equation(equation)
|
||||
return rendered is not None
|
||||
|
||||
|
||||
def process_candidate_file(candidate_file: str, pdfs_folder: str, sim_threshold: float) -> List[MathTest]:
|
||||
"""
|
||||
Process a single candidate file.
|
||||
Returns a list of MathTest objects extracted from the corresponding TeX file.
|
||||
"""
|
||||
logging.info("Processing %s", candidate_file)
|
||||
tests = []
|
||||
parse_result = parse_candidate_filename(candidate_file)
|
||||
if not parse_result:
|
||||
logging.error("Filename %s does not match expected format.", candidate_file)
|
||||
return tests
|
||||
|
||||
tex_basename, page_num = parse_result
|
||||
tex_file_path = os.path.join(pdfs_folder, f"{tex_basename}.tex")
|
||||
|
||||
if not os.path.exists(tex_file_path):
|
||||
logging.error("TeX file %s not found for candidate %s.", tex_file_path, candidate_file)
|
||||
return tests
|
||||
|
||||
candidate_text = extract_candidate_content(candidate_file)
|
||||
tex_content = extract_tex_content(tex_file_path)
|
||||
if not tex_content or not candidate_text or len(tex_content.strip()) < 100 or len(candidate_text.strip()) < 100:
|
||||
logging.error("No content extracted from %s", tex_file_path)
|
||||
return tests
|
||||
|
||||
matching_tex = find_matching_content(candidate_text, tex_content, sim_threshold)
|
||||
if not matching_tex:
|
||||
logging.warning("No matching TeX content found in %s for candidate %s", tex_file_path, candidate_file)
|
||||
return tests
|
||||
|
||||
logging.debug("Matching TeX content: %s", matching_tex)
|
||||
|
||||
math_equations = extract_math_from_tex(matching_tex)
|
||||
if not math_equations:
|
||||
logging.warning("No math equations found in matching content for candidate %s", candidate_file)
|
||||
return tests
|
||||
|
||||
# Filter out equations that are too short, remove duplicates, and shuffle
|
||||
math_equations = [(eq_type, eq.strip()) for (eq_type, eq) in math_equations if len(eq.strip()) > 20]
|
||||
math_equations = list(set(math_equations))
|
||||
random.shuffle(math_equations)
|
||||
|
||||
for i, (eq_type, equation) in enumerate(math_equations):
|
||||
if validate_equation(equation):
|
||||
test_id = f"{tex_basename}_pg{page_num}_math_{i:03d}"
|
||||
math_test = MathTest(
|
||||
id=test_id,
|
||||
pdf=f"{tex_basename}.pdf",
|
||||
page=page_num,
|
||||
type="math",
|
||||
math=equation,
|
||||
)
|
||||
tests.append(math_test)
|
||||
if len(tests) >= 10:
|
||||
break
|
||||
|
||||
return tests
|
||||
|
||||
|
||||
def process_tex_file_group(tex_basename: str, candidate_files: List[str], pdfs_folder: str, sim_threshold: float, max_pages: int) -> List[MathTest]:
|
||||
"""
|
||||
For a given TeX file, group candidate files by page, randomly shuffle the pages,
|
||||
and process them one-by-one. Stop once max_pages (pages with valid equations) have
|
||||
been processed.
|
||||
"""
|
||||
tests = []
|
||||
valid_pages = set()
|
||||
|
||||
# Group candidate files by page number.
|
||||
page_dict: Dict[int, List[str]] = {}
|
||||
for candidate_file in candidate_files:
|
||||
parse_result = parse_candidate_filename(candidate_file)
|
||||
if not parse_result:
|
||||
continue
|
||||
_, page_num = parse_result
|
||||
page_dict.setdefault(page_num, []).append(candidate_file)
|
||||
|
||||
# For each page, randomly choose one candidate file.
|
||||
distinct_candidate_files = []
|
||||
for page_num, files in page_dict.items():
|
||||
chosen_file = random.choice(files)
|
||||
distinct_candidate_files.append(chosen_file)
|
||||
|
||||
# Shuffle the pages randomly.
|
||||
random.shuffle(distinct_candidate_files)
|
||||
|
||||
# Process pages sequentially until max_pages with valid equations have been found.
|
||||
for candidate_file in distinct_candidate_files:
|
||||
result = process_candidate_file(candidate_file, pdfs_folder, sim_threshold)
|
||||
if result:
|
||||
tests.extend(result)
|
||||
# Mark this page as valid.
|
||||
page_num = parse_candidate_filename(candidate_file)[1]
|
||||
valid_pages.add(page_num)
|
||||
if len(valid_pages) >= max_pages:
|
||||
break
|
||||
|
||||
return tests
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Extract math equations from candidate files and corresponding TeX bases.")
|
||||
parser.add_argument("--math_data", required=True, help="Path to math_data folder")
|
||||
parser.add_argument("--candidate", required=True, help="Candidate folder name inside math_data")
|
||||
parser.add_argument("--max_pages", type=int, default=1, help="Maximum distinct pages with equations to process per TeX document")
|
||||
parser.add_argument("--parallel", type=int, default=8, help="Maximum process pool workers")
|
||||
parser.add_argument("--sim_threshold", type=float, default=0.7, help="Similarity threshold for matching candidate text")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
candidate_folder = os.path.join(args.math_data, args.candidate)
|
||||
pdfs_folder = os.path.join(args.math_data, "pdfs")
|
||||
|
||||
candidate_files = glob.glob(os.path.join(candidate_folder, "*.md"))
|
||||
logging.info("Found %d candidate files.", len(candidate_files))
|
||||
|
||||
# Group candidate files by TeX basename.
|
||||
tex_groups: Dict[str, List[str]] = {}
|
||||
for candidate_file in candidate_files:
|
||||
parse_result = parse_candidate_filename(candidate_file)
|
||||
if not parse_result:
|
||||
continue
|
||||
tex_basename, _ = parse_result
|
||||
tex_groups.setdefault(tex_basename, []).append(candidate_file)
|
||||
logging.info("Found %d TeX groups.", len(tex_groups))
|
||||
|
||||
# Remove output file if it exists to start fresh
|
||||
output_file = os.path.join(args.math_data, "math_tests.jsonl")
|
||||
if os.path.exists(output_file):
|
||||
os.remove(output_file)
|
||||
|
||||
all_math_tests = []
|
||||
|
||||
# Process each TeX group in parallel using ProcessPoolExecutor.
|
||||
with ProcessPoolExecutor(max_workers=args.parallel) as executor:
|
||||
future_to_tex = {
|
||||
executor.submit(process_tex_file_group, tex_basename, candidate_list, pdfs_folder, args.sim_threshold, args.max_pages): tex_basename
|
||||
for tex_basename, candidate_list in tex_groups.items()
|
||||
}
|
||||
for future in tqdm(as_completed(future_to_tex), total=len(future_to_tex), desc="Processing TeX files"):
|
||||
tex_basename = future_to_tex[future]
|
||||
try:
|
||||
tests = future.result()
|
||||
all_math_tests.extend(tests)
|
||||
# Incrementally save tests as each TeX group finishes processing.
|
||||
save_tests(all_math_tests, output_file)
|
||||
except Exception as e:
|
||||
logging.error("Error processing TeX group %s: %s", tex_basename, e)
|
||||
|
||||
logging.info("Found %d valid math equations from %d TeX groups.", len(all_math_tests), len(tex_groups))
|
||||
logging.info("Results incrementally saved to %s", output_file)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,405 @@
|
||||
"""
|
||||
mine_multi_column.py - Extract text from PDF documents which has multiple columns.
|
||||
|
||||
This script:
|
||||
1. Takes a file containing folder path which contains PDF documents as input
|
||||
2. Process each PDF to generate an HTML representation
|
||||
3. For each PDF, it renders to an image
|
||||
4. Uses Claude Sonnet to identify text from multiple columns in the rendered image
|
||||
5. Creates a test file asserting that the order (before/after) of text should appear
|
||||
6. Extracts the page from the PDF and saves it to an output folder
|
||||
|
||||
Usage:
|
||||
python mine_headers_footers.py --input_dir path/to/pdfs --output_dir path/to/output --api_key your_anthropic_api_key
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import subprocess
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Dict, List
|
||||
|
||||
import pypdf
|
||||
from anthropic import Anthropic
|
||||
from bs4 import BeautifulSoup
|
||||
from playwright.async_api import async_playwright
|
||||
from syntok.segmenter import process
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.data.renderpdf import (
|
||||
get_png_dimensions_from_base64,
|
||||
render_pdf_to_base64png,
|
||||
)
|
||||
|
||||
|
||||
def extract_code_block(initial_response):
|
||||
html_blocks = re.findall(r"```html\n(.*?)```", initial_response, re.DOTALL)
|
||||
if html_blocks:
|
||||
return html_blocks[-1].strip()
|
||||
code_blocks = re.findall(r"```\n(.*?)```", initial_response, re.DOTALL)
|
||||
if code_blocks:
|
||||
return code_blocks[-1].strip()
|
||||
html_blocks_no_newline = re.findall(r"```html(.*?)```", initial_response, re.DOTALL)
|
||||
if html_blocks_no_newline:
|
||||
return html_blocks_no_newline[-1].strip()
|
||||
code_blocks_no_newline = re.findall(r"```(.*?)```", initial_response, re.DOTALL)
|
||||
if code_blocks_no_newline:
|
||||
return code_blocks_no_newline[-1].strip()
|
||||
return None
|
||||
|
||||
|
||||
def generate_html_from_image(client, image_base64):
|
||||
"""Call Claude API to generate HTML from an image using a multi-step prompting strategy."""
|
||||
png_width, png_height = get_png_dimensions_from_base64(image_base64)
|
||||
try:
|
||||
analysis_response = client.messages.create(
|
||||
model="claude-3-7-sonnet-20250219",
|
||||
max_tokens=2000,
|
||||
temperature=0.1,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": image_base64}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"Analyze this document and provide a detailed assessment of its structure. "
|
||||
"Focus on the layout, headings, footers, and any complex formatting. Please be precise."
|
||||
),
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
analysis_text = ""
|
||||
for content in analysis_response.content:
|
||||
if content.type == "text":
|
||||
analysis_text += content.text
|
||||
|
||||
initial_response = client.messages.create(
|
||||
model="claude-3-7-sonnet-20250219",
|
||||
max_tokens=6000,
|
||||
temperature=0.2,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": image_base64}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"Render this document as clean, semantic HTML. Here is the analysis of the document structure:\n\n"
|
||||
f"{analysis_text}\n\n"
|
||||
"Requirements:\n"
|
||||
"1. Use appropriate HTML tags for headings, paragraphs, and lists.\n"
|
||||
"2. Use <header> and <footer> for top and bottom content.\n"
|
||||
"3. For images, use a placeholder <div> with class 'image'.\n"
|
||||
"4. Render math equations inline using \\( \\) or \\[ \\].\n"
|
||||
"5. Preserve any multi-column layout using CSS flexbox or grid.\n"
|
||||
f"6. The viewport is fixed at {png_width // 2}x{png_height // 2} pixels.\n\n"
|
||||
"Enclose your HTML in a ```html code block."
|
||||
),
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
initial_html = ""
|
||||
for content in initial_response.content:
|
||||
if content.type == "text":
|
||||
initial_html += content.text
|
||||
|
||||
return extract_code_block(initial_html)
|
||||
except Exception as e:
|
||||
print(f"Error calling Claude API: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def generate_tests_from_html(html_content: str, pdf_id: str, page_num: int) -> List[Dict]:
|
||||
"""
|
||||
Generate order tests from HTML content by splitting the main text into sentences and pairing them.
|
||||
Only order tests are generated.
|
||||
"""
|
||||
tests = []
|
||||
pdf_filename = pdf_id
|
||||
soup = BeautifulSoup(html_content, "html.parser")
|
||||
full_text = soup.get_text(separator=" ").strip()
|
||||
|
||||
sentences = []
|
||||
for paragraph in process(full_text):
|
||||
for sentence in paragraph:
|
||||
sentence_str = "".join(token.spacing + token.value for token in sentence).strip()
|
||||
if sentence_str:
|
||||
sentences.append(sentence_str)
|
||||
|
||||
if len(sentences) < 2:
|
||||
return tests
|
||||
|
||||
all_indexes = list(range(len(sentences)))
|
||||
random.shuffle(all_indexes)
|
||||
random_pairs = [(all_indexes[i * 2], all_indexes[i * 2 + 1]) for i in range(len(all_indexes) // 2)]
|
||||
random_pairs = [(min(i, j), max(i, j)) for (i, j) in random_pairs]
|
||||
|
||||
num_order_tests = 0
|
||||
for i, j in random_pairs:
|
||||
first_sentence = sentences[i]
|
||||
second_sentence = sentences[j]
|
||||
if len(first_sentence) < 10 or len(second_sentence) < 10:
|
||||
continue
|
||||
first_sentence = first_sentence.split("\n")[0].strip()
|
||||
second_sentence = second_sentence.split("\n")[0].strip()
|
||||
max_diffs = round(max(len(first_sentence), len(second_sentence)) * 0.05)
|
||||
if max_diffs > len(first_sentence) // 2 or max_diffs > len(second_sentence) // 2:
|
||||
continue
|
||||
tests.append(
|
||||
{
|
||||
"pdf": pdf_filename,
|
||||
"page": page_num,
|
||||
"id": f"{pdf_id}_order_{uuid.uuid4().hex[:8]}",
|
||||
"type": "order",
|
||||
"before": first_sentence,
|
||||
"after": second_sentence,
|
||||
"max_diffs": max_diffs,
|
||||
}
|
||||
)
|
||||
num_order_tests += 1
|
||||
if num_order_tests >= 5:
|
||||
break
|
||||
|
||||
return tests
|
||||
|
||||
|
||||
async def render_pdf_with_playwright(html_content, output_pdf_path, png_width, png_height):
|
||||
"""
|
||||
Render HTML content using Playwright and save it as a PDF.
|
||||
Tries different scale factors until the output PDF has exactly one page.
|
||||
"""
|
||||
scale_factors = [1.0, 0.9, 0.8, 0.7, 0.6, 0.5]
|
||||
for scale in scale_factors:
|
||||
try:
|
||||
async with async_playwright() as p:
|
||||
browser = await p.chromium.launch()
|
||||
page = await browser.new_page(viewport={"width": int(png_width // 2 * scale), "height": int(png_height // 2 * scale)})
|
||||
await page.set_content(html_content)
|
||||
|
||||
# Add KaTeX assets for math rendering
|
||||
katex_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "katex")
|
||||
katex_css_path = os.path.join(katex_dir, "katex.min.css")
|
||||
katex_js_path = os.path.join(katex_dir, "katex.min.js")
|
||||
katex_autorender_js_path = os.path.join(katex_dir, "auto-render.min.js")
|
||||
|
||||
await page.add_style_tag(path=katex_css_path)
|
||||
await page.add_script_tag(path=katex_js_path)
|
||||
await page.add_script_tag(path=katex_autorender_js_path)
|
||||
|
||||
await page.evaluate("""
|
||||
renderMathInElement(document.body, {
|
||||
delimiters: [
|
||||
{left: '\\\$begin:math:text$', right: '\\\\\\$end:math:text$', display: false},
|
||||
{left: '\\\$begin:math:display$', right: '\\\\\\$end:math:display$', display: true}
|
||||
],
|
||||
throwOnError: false
|
||||
});
|
||||
""")
|
||||
|
||||
await page.pdf(path=output_pdf_path, scale=scale, print_background=True)
|
||||
await browser.close()
|
||||
|
||||
try:
|
||||
reader = pypdf.PdfReader(output_pdf_path)
|
||||
if len(reader.pages) == 1:
|
||||
print(f"Successfully rendered as a single page PDF with scale factor {scale}")
|
||||
return True
|
||||
else:
|
||||
print(f"PDF has {len(reader.pages)} pages with scale factor {scale}, trying a smaller scale...")
|
||||
except Exception as pdf_check_error:
|
||||
print(f"Error checking PDF page count: {pdf_check_error}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error rendering PDF with Playwright at scale {scale}: {e}")
|
||||
print("Failed to render PDF as a single page with any scale factor")
|
||||
return False
|
||||
|
||||
|
||||
def process_pdf(pdf_info, args, client):
|
||||
"""
|
||||
Process a single PDF from a local folder:
|
||||
- Select a random page from the PDF.
|
||||
- Render that page as a base64 PNG.
|
||||
- Generate HTML from the image using Claude.
|
||||
- Optionally render the HTML to a PDF using Playwright.
|
||||
- Generate order tests from the HTML.
|
||||
"""
|
||||
local_pdf_path, index = pdf_info
|
||||
original_pdf_name = os.path.basename(local_pdf_path)
|
||||
pdf_id = original_pdf_name
|
||||
temp_pdf_dir = os.path.join(args.temp_dir, f"{os.path.splitext(original_pdf_name)[0]}")
|
||||
os.makedirs(temp_pdf_dir, exist_ok=True)
|
||||
|
||||
try:
|
||||
reader = pypdf.PdfReader(local_pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
if num_pages == 0:
|
||||
print(f"PDF has no pages: {local_pdf_path}")
|
||||
return None
|
||||
page_num = random.randint(1, num_pages)
|
||||
except Exception as e:
|
||||
print(f"Error reading {local_pdf_path}: {e}")
|
||||
return None
|
||||
|
||||
try:
|
||||
image_base64 = render_pdf_to_base64png(local_pdf_path, page_num, target_longest_image_dim=2048)
|
||||
except Exception as e:
|
||||
print(f"Error rendering page {page_num} from {local_pdf_path}: {e}")
|
||||
return None
|
||||
|
||||
try:
|
||||
html_content = generate_html_from_image(client, image_base64)
|
||||
if not html_content:
|
||||
print(f"Failed to generate HTML for {local_pdf_path}, page {page_num}")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"Error generating HTML for {local_pdf_path}: {e}")
|
||||
return None
|
||||
|
||||
html_dir = os.path.join(args.output_dir, "html")
|
||||
pdfs_dir = os.path.join(args.output_dir, "pdfs")
|
||||
os.makedirs(html_dir, exist_ok=True)
|
||||
os.makedirs(pdfs_dir, exist_ok=True)
|
||||
|
||||
html_path = os.path.join(html_dir, f"{os.path.splitext(original_pdf_name)[0]}_page{page_num}.html")
|
||||
try:
|
||||
with open(html_path, "w") as f:
|
||||
f.write(html_content)
|
||||
except Exception as e:
|
||||
print(f"Error saving HTML for {local_pdf_path}: {e}")
|
||||
return None
|
||||
|
||||
playwright_pdf_path = None
|
||||
render_success = False
|
||||
playwright_pdf_filename = f"{os.path.splitext(original_pdf_name)[0]}_page{page_num}.pdf"
|
||||
if not args.skip_playwright:
|
||||
playwright_pdf_path = os.path.join(pdfs_dir, playwright_pdf_filename)
|
||||
try:
|
||||
png_width, png_height = get_png_dimensions_from_base64(image_base64)
|
||||
render_success = asyncio.run(render_pdf_with_playwright(html_content, playwright_pdf_path, png_width, png_height))
|
||||
if render_success:
|
||||
print(f"Successfully rendered with Playwright: {playwright_pdf_path}")
|
||||
else:
|
||||
print(f"Failed to render as a single page PDF: {playwright_pdf_path}")
|
||||
playwright_pdf_path = None
|
||||
except Exception as e:
|
||||
print(f"Failed to render with Playwright: {e}")
|
||||
playwright_pdf_path = None
|
||||
render_success = False
|
||||
if not args.skip_playwright and not render_success:
|
||||
return None
|
||||
|
||||
tests = generate_tests_from_html(html_content, pdf_id, page_num)
|
||||
# IMPORTANT: Preserve the original PDF filename in the tests.
|
||||
for test in tests:
|
||||
test["pdf"] = original_pdf_name
|
||||
|
||||
try:
|
||||
if os.path.exists(temp_pdf_dir):
|
||||
subprocess.run(["rm", "-rf", temp_pdf_dir])
|
||||
except Exception as e:
|
||||
print(f"Error cleaning up temp directory {temp_pdf_dir}: {e}")
|
||||
|
||||
return {
|
||||
"pdf_id": pdf_id,
|
||||
"pdf_path": local_pdf_path,
|
||||
"page_number": page_num,
|
||||
"html_path": html_path,
|
||||
"playwright_pdf_path": playwright_pdf_path,
|
||||
"tests": tests,
|
||||
"num_tests": len(tests),
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Convert PDFs in a folder to HTML templates and render with Playwright (order tests only)")
|
||||
parser.add_argument("--input_dir", required=True, help="Folder containing PDF files")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to store HTML and tests")
|
||||
parser.add_argument("--temp_dir", default="/tmp/mine_tables", help="Directory for temporary files")
|
||||
parser.add_argument("--max_tests", type=int, default=100, help="Maximum number of PDFs to process (randomly selected)")
|
||||
parser.add_argument("--parallel", type=int, default=1, help="Number of parallel threads to use")
|
||||
parser.add_argument("--api_key", help="Claude API key (or set ANTHROPIC_API_KEY environment variable)")
|
||||
parser.add_argument("--skip_playwright", action="store_true", help="Skip Playwright PDF rendering")
|
||||
parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
|
||||
args = parser.parse_args()
|
||||
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
os.makedirs(args.temp_dir, exist_ok=True)
|
||||
|
||||
api_key = args.api_key or os.environ.get("ANTHROPIC_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: API key not provided. Use --api_key or set ANTHROPIC_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
client = Anthropic(api_key=api_key)
|
||||
|
||||
pdf_paths = []
|
||||
for root, dirs, files in os.walk(args.input_dir):
|
||||
for file in files:
|
||||
if file.lower().endswith(".pdf"):
|
||||
pdf_paths.append(os.path.join(root, file))
|
||||
print(f"Found {len(pdf_paths)} PDF files in {args.input_dir}")
|
||||
|
||||
random.shuffle(pdf_paths)
|
||||
pdf_paths = pdf_paths[: args.max_tests]
|
||||
|
||||
synthetic_json_path = os.path.join(args.output_dir, "synthetic.jsonl")
|
||||
open(synthetic_json_path, "w").close()
|
||||
|
||||
test_counter = 0
|
||||
test_types = {"order": 0}
|
||||
results = []
|
||||
|
||||
import threading
|
||||
|
||||
file_lock = threading.Lock()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=args.parallel) as executor:
|
||||
futures = {executor.submit(process_pdf, (pdf_path, i), args, client): pdf_path for i, pdf_path in enumerate(pdf_paths)}
|
||||
for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Processing PDFs"):
|
||||
pdf_path = futures[future]
|
||||
try:
|
||||
result = future.result()
|
||||
if result and result.get("tests"):
|
||||
results.append(result)
|
||||
with file_lock:
|
||||
with open(synthetic_json_path, "a") as f:
|
||||
for test in result["tests"]:
|
||||
f.write(json.dumps(test) + "\n")
|
||||
test_counter += len(result["tests"])
|
||||
for test in result["tests"]:
|
||||
if test.get("type") == "order":
|
||||
test_types["order"] += 1
|
||||
print(f"Added {len(result['tests'])} tests from {result['pdf_id']}, total: {test_counter}")
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_path}: {e}")
|
||||
|
||||
print(f"Generated {len(results)} HTML templates")
|
||||
if not args.skip_playwright:
|
||||
playwright_success = sum(1 for r in results if r and r.get("playwright_pdf_path"))
|
||||
print(f"Playwright PDF rendering: {playwright_success}/{len(results)} successful")
|
||||
print(f"Saved {test_counter} tests to {synthetic_json_path}")
|
||||
print(f"Generated a total of {test_counter} tests across {len(results)} templates")
|
||||
print("Test type distribution:")
|
||||
for test_type, count in test_types.items():
|
||||
print(f" - {test_type}: {count} tests")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,295 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_multilingual_gpt.py - Identify PDF documents with tables and copy them.
|
||||
|
||||
This script:
|
||||
1. Takes a file containing S3 paths to PDF documents as input
|
||||
2. For each PDF, renders a random page and uses GPT-4o to check the primary language
|
||||
3. Identifies PDFs where the page is not in english
|
||||
4. Copies those PDF files to a new output folder
|
||||
|
||||
Usage:
|
||||
python mine_multilingual_gpt.py --input_list path/to/s3_paths.txt --output_dir path/to/output --api_key your_openai_api_key
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import random
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Optional
|
||||
|
||||
import boto3
|
||||
import pypdf
|
||||
from lingua import Language
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
|
||||
TARGET_IMAGE_DIM = 1024
|
||||
|
||||
|
||||
class LanguageDetectionResponse(BaseModel):
|
||||
"""Structured output for table detection."""
|
||||
|
||||
two_letter_language_code: str
|
||||
|
||||
|
||||
def download_pdf_from_s3(s3_path: str, local_path: str) -> bool:
|
||||
"""
|
||||
Download a PDF file from S3.
|
||||
|
||||
Args:
|
||||
s3_path: The S3 path (s3://bucket/path/to/file.pdf)
|
||||
local_path: The local path to save the file
|
||||
|
||||
Returns:
|
||||
bool: True if download was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Parse S3 path
|
||||
parts = s3_path.replace("s3://", "").split("/", 1)
|
||||
bucket = parts[0]
|
||||
key = parts[1]
|
||||
|
||||
# Create S3 client
|
||||
s3 = boto3.client("s3")
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
|
||||
# Download file
|
||||
s3.download_file(bucket, key, local_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {s3_path}: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def check_for_table(pdf_path: str, page_num: int, api_key: str) -> Optional[bool]:
|
||||
"""
|
||||
Use GPT-4o to check if a page contains a table.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
Optional[bool]: True if page contains a table, False otherwise, None if detection fails
|
||||
"""
|
||||
# Initialize OpenAI client
|
||||
client = OpenAI(api_key=api_key)
|
||||
|
||||
try:
|
||||
# Render the PDF page as an image (render_pdf_to_base64png is 1-indexed)
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=TARGET_IMAGE_DIM)
|
||||
|
||||
# Simple prompt asking about tables
|
||||
prompt = "Respond with the primary language that this document is written in. Reply in JSON, and use a two letter ISO 639 language code, such as 'en', 'pt', 'fr', etc."
|
||||
|
||||
response = client.beta.chat.completions.parse(
|
||||
model="gpt-4o-2024-08-06",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}}],
|
||||
}
|
||||
],
|
||||
temperature=0.1,
|
||||
max_tokens=100,
|
||||
response_format=LanguageDetectionResponse,
|
||||
)
|
||||
|
||||
if not response.choices or len(response.choices) == 0:
|
||||
print(f"No response generated for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
# Parse the structured response
|
||||
parsed_response = response.choices[0].message.parsed
|
||||
|
||||
if parsed_response is None:
|
||||
print(f"Failed to parse response for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
is_non_english = parsed_response.two_letter_language_code.lower() != "en"
|
||||
|
||||
if is_non_english:
|
||||
print(f"Found {parsed_response.two_letter_language_code} language in {pdf_path} page {page_num + 1}")
|
||||
|
||||
return is_non_english
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error checking {pdf_path} page {page_num}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def process_pdf(s3_path: str, temp_dir: str, output_dir: str, api_key: str) -> bool:
|
||||
"""
|
||||
Process a single PDF from S3.
|
||||
|
||||
Args:
|
||||
s3_path: S3 path to the PDF
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
bool: True if the PDF has a table and was copied, False otherwise
|
||||
"""
|
||||
# Extract filename from S3 path
|
||||
pdf_filename = os.path.basename(s3_path)
|
||||
local_pdf_path = os.path.join(temp_dir, pdf_filename)
|
||||
|
||||
# Download PDF from S3
|
||||
if not download_pdf_from_s3(s3_path, local_pdf_path):
|
||||
return False
|
||||
|
||||
pdf_filter = PdfFilter(languages_to_keep=Language.all())
|
||||
|
||||
if pdf_filter.filter_out_pdf(local_pdf_path):
|
||||
print(f"Filtering out {pdf_filename}")
|
||||
return False
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(local_pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return False
|
||||
|
||||
# Select a random page to check
|
||||
page_num = random.randint(0, num_pages - 1)
|
||||
page_num = random.choice([page_num, 0]) # Bias 50% of the time to do the first page
|
||||
|
||||
# Check if the page contains a table
|
||||
has_table = check_for_table(local_pdf_path, page_num, api_key)
|
||||
|
||||
if has_table:
|
||||
# Extract just the page with the table and save it as a new PDF
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Create output filename with basename_pgnum.pdf format
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(output_dir, f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
|
||||
# Extract the single page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_pdf_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
print(f"Extracted page {page_num+1} with table from {pdf_filename} to {os.path.basename(output_pdf_path)}")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
return False
|
||||
finally:
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Identify and copy PDFs with tables")
|
||||
parser.add_argument("--input_list", required=True, help="Path to a file containing S3 paths to PDFs")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to copy PDFs with tables")
|
||||
parser.add_argument("--api_key", help="OpenAI API key (if not provided, will use OPENAI_API_KEY environment variable)")
|
||||
parser.add_argument("--temp_dir", default="/tmp/mine_tables", help="Directory for temporary files")
|
||||
parser.add_argument("--max_pdfs", type=int, default=100, help="Maximum number of PDFs with tables to find")
|
||||
parser.add_argument("--parallel", type=int, default=1, help="Number of parallel workers (default: 1 for sequential)")
|
||||
parser.add_argument("--reservoir_multiplier", type=int, default=100, help="Multiplier for reservoir sampling (default: 100x max_pdfs)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: OpenAI API key not provided. Use --api_key or set OPENAI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
os.makedirs(args.temp_dir, exist_ok=True)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Reservoir sampling to get random subset of PDFs
|
||||
reservoir_size = args.max_pdfs * args.reservoir_multiplier
|
||||
pdf_paths = []
|
||||
n = 0 # Total number of items seen
|
||||
|
||||
print(f"Using reservoir sampling with size {reservoir_size}")
|
||||
|
||||
with open(args.input_list, "r") as f:
|
||||
for line in f:
|
||||
n += 1
|
||||
path = line.strip()
|
||||
if not path:
|
||||
continue
|
||||
|
||||
if len(pdf_paths) < reservoir_size:
|
||||
pdf_paths.append(path)
|
||||
else:
|
||||
# Randomly decide whether to include this item
|
||||
s = random.randint(1, n)
|
||||
if s <= reservoir_size:
|
||||
pdf_paths[s - 1] = path
|
||||
|
||||
# Shuffle the reservoir
|
||||
random.shuffle(pdf_paths)
|
||||
|
||||
print(f"Sampled {len(pdf_paths)} PDF paths from {n} total paths")
|
||||
|
||||
table_pdfs_found = 0
|
||||
|
||||
if args.parallel > 1:
|
||||
# Parallel processing
|
||||
print(f"Processing PDFs with {args.parallel} parallel workers")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=args.parallel) as executor:
|
||||
futures = []
|
||||
|
||||
# Submit all tasks
|
||||
for s3_path in pdf_paths:
|
||||
if table_pdfs_found >= args.max_pdfs:
|
||||
break
|
||||
future = executor.submit(process_pdf, s3_path, args.temp_dir, args.output_dir, api_key)
|
||||
futures.append(future)
|
||||
|
||||
# Process results as they complete
|
||||
with tqdm(total=min(len(pdf_paths), args.max_pdfs), desc="Processing PDFs") as pbar:
|
||||
for future in as_completed(futures):
|
||||
try:
|
||||
result = future.result()
|
||||
if result:
|
||||
table_pdfs_found += 1
|
||||
pbar.update(1)
|
||||
|
||||
if table_pdfs_found >= args.max_pdfs:
|
||||
print(f"Reached maximum number of PDFs with tables ({args.max_pdfs}), stopping")
|
||||
# Cancel remaining futures
|
||||
for f in futures:
|
||||
f.cancel()
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"Error in parallel processing: {str(e)}")
|
||||
else:
|
||||
# Sequential processing
|
||||
for s3_path in tqdm(pdf_paths, desc="Processing PDFs"):
|
||||
if process_pdf(s3_path, args.temp_dir, args.output_dir, api_key):
|
||||
table_pdfs_found += 1
|
||||
|
||||
if table_pdfs_found >= args.max_pdfs:
|
||||
print(f"Reached maximum number of PDFs with tables ({args.max_pdfs}), stopping")
|
||||
break
|
||||
|
||||
print(f"Found and copied {table_pdfs_found} PDFs to {args.output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,120 @@
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
|
||||
def process_jsonl(jsonl_file, output_dir="downloaded_images"):
|
||||
"""
|
||||
Process each line in the JSONL file and download the corresponding images.
|
||||
|
||||
Args:
|
||||
jsonl_file (str): Path to the JSONL file
|
||||
output_dir (str): Directory to save downloaded images
|
||||
"""
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
with open(jsonl_file, "r") as f:
|
||||
for line_num, line in enumerate(f, 1):
|
||||
try:
|
||||
data = json.loads(line.strip())
|
||||
pdf_path = data.get("pdf", "")
|
||||
pdf_filename = os.path.basename(pdf_path)
|
||||
url = data.get("url", "")
|
||||
if not url:
|
||||
print(f"Line {line_num}: Missing URL, skipping...")
|
||||
continue
|
||||
|
||||
print(f"Processing line {line_num}: {pdf_filename} - {url}")
|
||||
image_path = download_image(url, pdf_filename, output_dir)
|
||||
|
||||
if image_path:
|
||||
print(f"Successfully downloaded: {image_path}")
|
||||
else:
|
||||
print(f"Failed to download image for {pdf_filename}")
|
||||
time.sleep(1)
|
||||
|
||||
except json.JSONDecodeError:
|
||||
print(f"Line {line_num}: Invalid JSON, skipping...")
|
||||
except Exception as e:
|
||||
print(f"Line {line_num}: Error processing - {str(e)}")
|
||||
|
||||
|
||||
def download_image(url, output_filename, output_dir):
|
||||
"""
|
||||
Download the highest resolution JPEG (before JPEG2000) from the Library of Congress URL.
|
||||
|
||||
Args:
|
||||
url (str): The initial URL from the JSONL file
|
||||
output_filename (str): Filename for the downloaded image
|
||||
output_dir (str): Directory to save the image
|
||||
|
||||
Returns:
|
||||
str or None: Path to the downloaded image if successful, None otherwise
|
||||
"""
|
||||
try:
|
||||
response = requests.get(url)
|
||||
response.raise_for_status()
|
||||
soup = BeautifulSoup(response.text, "html.parser")
|
||||
loc_link = soup.find("a", class_="btn btn-outline-primary text-nowrap", title="View the original source for this item in a new tab")
|
||||
if not loc_link:
|
||||
print(f"Could not find 'View on www.loc.gov' link on {url}")
|
||||
return None
|
||||
|
||||
loc_url = loc_link["href"]
|
||||
print(f"Found LOC URL: {loc_url}")
|
||||
response = requests.get(loc_url)
|
||||
response.raise_for_status()
|
||||
soup = BeautifulSoup(response.text, "html.parser")
|
||||
download_select = soup.find("select", id="download")
|
||||
if not download_select:
|
||||
print(f"Could not find download options on {loc_url}")
|
||||
return None
|
||||
|
||||
jpeg_options = [option for option in download_select.find_all("option") if "JPEG" in option.text and "JPEG2000" not in option.text]
|
||||
|
||||
if not jpeg_options:
|
||||
print(f"No JPEG options found on {loc_url}")
|
||||
return None
|
||||
|
||||
highest_jpeg = jpeg_options[-1]
|
||||
image_url = highest_jpeg["value"]
|
||||
|
||||
print(f"Found highest resolution JPEG: {image_url}")
|
||||
response = requests.get(image_url, stream=True)
|
||||
response.raise_for_status()
|
||||
|
||||
if not output_filename.lower().endswith((".jpg", ".jpeg")):
|
||||
output_filename = f"{os.path.splitext(output_filename)[0]}.jpg"
|
||||
|
||||
output_path = os.path.join(output_dir, output_filename)
|
||||
|
||||
with open(output_path, "wb") as f:
|
||||
for chunk in response.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
|
||||
return output_path
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error downloading image: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def main():
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Download images from Library of Congress based on JSONL file")
|
||||
parser.add_argument("jsonl_file", help="Path to the JSONL file containing URLs")
|
||||
parser.add_argument("--output", "-o", default="downloaded_images", help="Directory to save downloaded images (default: downloaded_images)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"Processing JSONL file: {args.jsonl_file}")
|
||||
process_jsonl(args.jsonl_file, args.output)
|
||||
print("Processing complete!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,122 @@
|
||||
import json
|
||||
import random
|
||||
import re
|
||||
|
||||
|
||||
def extract_random_segment(text, min_words=7, max_words=15):
|
||||
"""Extract a random segment of 7-15 words from the text."""
|
||||
words = text.split()
|
||||
if len(words) <= max_words:
|
||||
return text # Return full text if it's shorter than max_words
|
||||
|
||||
max_start = len(words) - min_words
|
||||
start = random.randint(0, max_start)
|
||||
remaining_words = len(words) - start
|
||||
segment_length = random.randint(min_words, min(max_words, remaining_words))
|
||||
segment = words[start : start + segment_length]
|
||||
return " ".join(segment)
|
||||
|
||||
|
||||
def process_jsonl_file_present(input_file, output_file):
|
||||
"""Process a JSONL file and create multiple random cases for each PDF."""
|
||||
with open(input_file, "r") as infile, open(output_file, "w") as outfile:
|
||||
for line in infile:
|
||||
if line.strip(): # Skip empty lines
|
||||
data = json.loads(line)
|
||||
image = data["image"]
|
||||
original_text = data["text"]
|
||||
num_cases = random.randint(1, 3)
|
||||
|
||||
for _ in range(num_cases):
|
||||
processed_num = random.randint(5, 10)
|
||||
processed_id = f"{image}_processed{processed_num:02d}"
|
||||
max_diffs = random.randint(1, 2)
|
||||
text_segment = extract_random_segment(original_text)
|
||||
|
||||
new_case = {
|
||||
"pdf": f"{image}.pdf",
|
||||
"page": 1,
|
||||
"id": processed_id,
|
||||
"type": "present",
|
||||
"max_diffs": max_diffs,
|
||||
"text": text_segment,
|
||||
"case_sensitive": True,
|
||||
"first_n": None,
|
||||
"last_n": None,
|
||||
}
|
||||
outfile.write(json.dumps(new_case) + "\n")
|
||||
|
||||
|
||||
def extract_ordered_segments(text, min_words=7, max_words=15):
|
||||
"""Extract two ordered segments from the text."""
|
||||
sentences = re.split(r"(?<=[.!?])\s+", text)
|
||||
|
||||
if len(sentences) < 2:
|
||||
return None, None
|
||||
valid_indices = list(range(len(sentences)))
|
||||
if len(valid_indices) <= 2:
|
||||
before_idx, after_idx = 0, 1
|
||||
else:
|
||||
before_idx = random.randint(0, len(valid_indices) - 2)
|
||||
after_idx = random.randint(before_idx + 1, len(valid_indices) - 1)
|
||||
before_sentence = sentences[before_idx]
|
||||
after_sentence = sentences[after_idx]
|
||||
|
||||
before_words = before_sentence.split()
|
||||
after_words = after_sentence.split()
|
||||
|
||||
if len(before_words) > max_words:
|
||||
start = random.randint(0, len(before_words) - min_words)
|
||||
length = random.randint(min_words, min(max_words, len(before_words) - start))
|
||||
before_segment = " ".join(before_words[start : start + length])
|
||||
else:
|
||||
before_segment = before_sentence
|
||||
|
||||
if len(after_words) > max_words:
|
||||
start = random.randint(0, len(after_words) - min_words)
|
||||
length = random.randint(min_words, min(max_words, len(after_words) - start))
|
||||
after_segment = " ".join(after_words[start : start + length])
|
||||
else:
|
||||
after_segment = after_sentence
|
||||
|
||||
return before_segment, after_segment
|
||||
|
||||
|
||||
def process_jsonl_file_order(input_file, output_file):
|
||||
"""Process a JSONL file and create order-type cases."""
|
||||
with open(input_file, "r") as infile, open(output_file, "w") as outfile:
|
||||
for line in infile:
|
||||
if line.strip(): # Skip empty lines
|
||||
data = json.loads(line)
|
||||
image = data["image"]
|
||||
original_text = data["text"]
|
||||
num_cases = random.randint(1, 3)
|
||||
|
||||
for _ in range(num_cases):
|
||||
before_text, after_text = extract_ordered_segments(original_text)
|
||||
if not before_text or not after_text:
|
||||
continue
|
||||
processed_num = random.randint(11, 16)
|
||||
processed_id = f"{image}_processed{processed_num:02d}"
|
||||
max_diffs = random.randint(1, 3)
|
||||
|
||||
new_case = {
|
||||
"pdf": f"{image}.pdf",
|
||||
"page": 1,
|
||||
"id": processed_id,
|
||||
"type": "order",
|
||||
"before": before_text,
|
||||
"after": after_text,
|
||||
"max_diffs": max_diffs,
|
||||
"checked": "verified",
|
||||
"url": f"https://example.com/document/{image}",
|
||||
}
|
||||
|
||||
outfile.write(json.dumps(new_case) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
input_file = "olmoce/bench/sample_data/old_scans.jsonl"
|
||||
output_file = "order_cases.jsonl"
|
||||
process_jsonl_file_present(input_file, output_file)
|
||||
process_jsonl_file_order(input_file, output_file)
|
||||
@@ -0,0 +1,277 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_old_scans_math.py - Extract mathematical equations from PDF documents.
|
||||
|
||||
This script:
|
||||
1. Takes a folder containing PDF documents as input
|
||||
2. For each PDF, extracts a random page and renders it to an image
|
||||
3. Uses Gemini to identify mathematical equations in the rendered image
|
||||
4. Creates a test file asserting that the equation text should be present
|
||||
5. Extracts the page from the PDF and saves it to an output folder
|
||||
|
||||
Usage:
|
||||
python mine_old_scans_math.py --input_dir path/to/pdf_folder --output_dir path/to/output --api_key your_gemini_api_key
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
from typing import List, Optional
|
||||
|
||||
import pypdf
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.bench.tests import TextPresenceTest, save_tests
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
|
||||
|
||||
def extract_page_from_pdf(input_path: str, output_path: str, page_num: int) -> bool:
|
||||
"""
|
||||
Extract a specific page from a PDF and save it as a new PDF.
|
||||
|
||||
Args:
|
||||
input_path: Path to the input PDF
|
||||
output_path: Path to save the extracted page
|
||||
page_num: The page number to extract (0-indexed)
|
||||
|
||||
Returns:
|
||||
bool: True if extraction was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Ensure output directory exists
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
# Read the input PDF
|
||||
reader = pypdf.PdfReader(input_path)
|
||||
|
||||
# Check if page number is valid
|
||||
if page_num >= len(reader.pages):
|
||||
print(f"Page number {page_num} out of range for {input_path} with {len(reader.pages)} pages")
|
||||
return False
|
||||
|
||||
# Create a new PDF with just the selected page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error extracting page {page_num} from {input_path}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def detect_equations(pdf_path: str, page_num: int, api_key: str) -> Optional[List[str]]:
|
||||
"""
|
||||
Use Gemini to detect mathematical equations in a rendered PDF page.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: Gemini API key
|
||||
|
||||
Returns:
|
||||
Optional[List[str]]: List of detected equations, or None if detection failed
|
||||
"""
|
||||
client = genai.Client(
|
||||
api_key=os.environ.get("GEMINI_API_KEY"),
|
||||
)
|
||||
model = "gemini-2.0-flash"
|
||||
|
||||
# Render the PDF page as an image
|
||||
try:
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=2048) # render_pdf_to_base64png is 1-indexed
|
||||
except Exception as e:
|
||||
print(f"Error rendering PDF page: {str(e)}")
|
||||
return None
|
||||
|
||||
image_part = types.Part(inline_data=types.Blob(mime_type="image/png", data=base64.b64decode(image_base64)))
|
||||
|
||||
contents = [
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
image_part,
|
||||
types.Part.from_text(
|
||||
text="""Please extract the mathematical equations from the document without omission. Always output the mathematical equations as Latex escaped with $$. Do not hallucinate"""
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
generate_content_config = types.GenerateContentConfig(
|
||||
temperature=1,
|
||||
top_p=0.95,
|
||||
top_k=40,
|
||||
max_output_tokens=8192,
|
||||
response_mime_type="application/json",
|
||||
response_schema=genai.types.Schema(
|
||||
type=genai.types.Type.OBJECT,
|
||||
properties={
|
||||
"equations": genai.types.Schema(
|
||||
type=genai.types.Type.ARRAY,
|
||||
items=genai.types.Schema(
|
||||
type=genai.types.Type.STRING,
|
||||
),
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
response = client.models.generate_content(model=model, contents=contents, config=generate_content_config)
|
||||
|
||||
assert len(response.candidates) > 0, "No candidates found"
|
||||
assert response.candidates[0].finish_reason == types.FinishReason.STOP, "Finish reason was not STOP, likely a processing error or repetition failure"
|
||||
|
||||
data = json.loads(response.candidates[0].content.parts[0].text)
|
||||
|
||||
return data.get("equations", [])
|
||||
|
||||
|
||||
def process_pdf(pdf_path: str, output_dir: str, api_key: str, tests: List[TextPresenceTest], max_pages_per_pdf: int = 10) -> None:
|
||||
"""
|
||||
Process a single PDF, extracting equations from multiple pages.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF
|
||||
output_dir: Directory for output files
|
||||
api_key: Gemini API key
|
||||
tests: List to append tests to
|
||||
max_pages_per_pdf: Maximum number of pages to process per PDF
|
||||
"""
|
||||
# Extract filename from path
|
||||
pdf_filename = os.path.basename(pdf_path)
|
||||
|
||||
pdf_filter = PdfFilter()
|
||||
|
||||
if pdf_filter.filter_out_pdf(pdf_path):
|
||||
print("Filtering out", pdf_filename)
|
||||
return
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return
|
||||
|
||||
# Get all pages and shuffle them to select a random subset
|
||||
all_pages = list(range(num_pages))
|
||||
random.shuffle(all_pages)
|
||||
|
||||
# Take only the specified maximum number of pages
|
||||
pages_to_process = all_pages[: min(max_pages_per_pdf, num_pages)]
|
||||
|
||||
processed_pages = 0
|
||||
|
||||
for page_num in pages_to_process:
|
||||
# Detect equations
|
||||
equations = detect_equations(pdf_path, page_num, api_key)
|
||||
|
||||
# Only keep equations that are non-empty
|
||||
equations = [eq for eq in equations if len(eq.strip()) > 3]
|
||||
|
||||
if not equations:
|
||||
print(f"No equations detected in {pdf_filename} page {page_num+1}")
|
||||
continue
|
||||
|
||||
# Extract the page and save to output dir
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(output_dir, "pdfs", f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
|
||||
extract_page_from_pdf(pdf_path, output_pdf_path, page_num)
|
||||
|
||||
# Create tests for each equation
|
||||
for i, equation in enumerate(equations):
|
||||
test_id = f"{pdf_basename}_pg{page_num+1}_equation_{i:02d}"
|
||||
test = TextPresenceTest(
|
||||
id=test_id,
|
||||
pdf=f"{pdf_basename}_pg{page_num+1}.pdf",
|
||||
page=1, # The extracted PDF has only one page
|
||||
type="present",
|
||||
text=equation,
|
||||
max_diffs=0,
|
||||
)
|
||||
tests.append(test)
|
||||
|
||||
print(f"Processed {pdf_filename} page {page_num+1}, found {len(equations)} equations")
|
||||
processed_pages += 1
|
||||
|
||||
print(f"Completed processing {processed_pages} pages from {pdf_filename}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
|
||||
|
||||
def get_pdf_files_from_directory(directory: str) -> List[str]:
|
||||
"""
|
||||
Get a list of all PDF files in a directory.
|
||||
|
||||
Args:
|
||||
directory: Path to the directory containing PDFs
|
||||
|
||||
Returns:
|
||||
List[str]: List of full paths to PDF files
|
||||
"""
|
||||
pdf_files = []
|
||||
|
||||
for filename in os.listdir(directory):
|
||||
if filename.lower().endswith(".pdf"):
|
||||
full_path = os.path.join(directory, filename)
|
||||
if os.path.isfile(full_path):
|
||||
pdf_files.append(full_path)
|
||||
|
||||
return pdf_files
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Extract mathematical equations from PDF documents")
|
||||
parser.add_argument("--input_dir", required=True, help="Directory containing PDF files")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to store extracted pages and tests")
|
||||
parser.add_argument("--api_key", help="Gemini API key (if not provided, will use GEMINI_API_KEY environment variable)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("GEMINI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: Gemini API key not provided. Use --api_key or set GEMINI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
# Create directories
|
||||
os.makedirs(os.path.join(args.output_dir, "pdfs"), exist_ok=True)
|
||||
|
||||
# Get PDF files from input directory
|
||||
pdf_files = get_pdf_files_from_directory(args.input_dir)
|
||||
|
||||
if not pdf_files:
|
||||
print(f"No PDF files found in {args.input_dir}")
|
||||
return
|
||||
|
||||
print(f"Found {len(pdf_files)} PDF files in input directory")
|
||||
|
||||
# Process each PDF
|
||||
tests = []
|
||||
for pdf_path in tqdm(pdf_files, desc="Processing PDFs"):
|
||||
process_pdf(pdf_path, args.output_dir, api_key, tests)
|
||||
|
||||
# Save tests after each PDF to avoid losing data in case of crashes
|
||||
if tests:
|
||||
save_tests(tests, os.path.join(args.output_dir, "equation_tests.jsonl"))
|
||||
|
||||
# if len(tests) > 100:
|
||||
# break
|
||||
|
||||
print(f"Saved {len(tests)} tests to {os.path.join(args.output_dir, 'equation_tests.jsonl')}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,530 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
analyze_documents.py - Analyze document layout and extract content from PDF documents.
|
||||
|
||||
This script:
|
||||
1. Takes a file containing S3 paths to PDF documents as input
|
||||
2. For each PDF, extracts a random page and renders it to an image
|
||||
3. Uses Gemini to analyze document layout features (columns, articles, text inserts, etc.)
|
||||
4. If specific layout features are detected, proceeds with full document content extraction
|
||||
5. Extracts the page from the PDF and saves it to an output folder along with analysis results
|
||||
|
||||
Usage:
|
||||
python analyze_documents.py --input_list path/to/s3_paths.txt --output_dir path/to/output --api_key your_gemini_api_key [--parallel 4]
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import threading
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import boto3
|
||||
import pypdf
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
|
||||
# Create a thread-safe lock for writing to output files
|
||||
file_lock = threading.Lock()
|
||||
|
||||
|
||||
def download_pdf_from_s3(s3_path: str, local_path: str) -> bool:
|
||||
"""
|
||||
Download a PDF file from S3.
|
||||
|
||||
Args:
|
||||
s3_path: The S3 path (s3://bucket/path/to/file.pdf)
|
||||
local_path: The local path to save the file
|
||||
|
||||
Returns:
|
||||
bool: True if download was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Parse S3 path
|
||||
parts = s3_path.replace("s3://", "").split("/", 1)
|
||||
bucket = parts[0]
|
||||
key = parts[1]
|
||||
|
||||
# Create S3 client
|
||||
s3 = boto3.client("s3")
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
|
||||
# Download file
|
||||
s3.download_file(bucket, key, local_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {s3_path}: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def extract_page_from_pdf(input_path: str, output_path: str, page_num: int) -> bool:
|
||||
"""
|
||||
Extract a specific page from a PDF and save it as a new PDF.
|
||||
|
||||
Args:
|
||||
input_path: Path to the input PDF
|
||||
output_path: Path to save the extracted page
|
||||
page_num: The page number to extract (0-indexed)
|
||||
|
||||
Returns:
|
||||
bool: True if extraction was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Ensure output directory exists
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
# Read the input PDF
|
||||
reader = pypdf.PdfReader(input_path)
|
||||
|
||||
# Check if page number is valid
|
||||
if page_num >= len(reader.pages):
|
||||
print(f"Page number {page_num} out of range for {input_path} with {len(reader.pages)} pages")
|
||||
return False
|
||||
|
||||
# Create a new PDF with just the selected page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error extracting page {page_num} from {input_path}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def analyze_document_layout(pdf_path: str, page_num: int, api_key: str) -> Optional[Tuple[Dict[str, Any], str]]:
|
||||
"""
|
||||
Use Gemini to analyze document layout features in a rendered PDF page.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: Gemini API key
|
||||
|
||||
Returns:
|
||||
Optional[Tuple[Dict[str, Any], str]]:
|
||||
A tuple with the layout analysis results as a dictionary and the base64 string of the rendered page image.
|
||||
Returns None if analysis fails.
|
||||
"""
|
||||
# Initialize Gemini client
|
||||
client = genai.Client(
|
||||
api_key=api_key,
|
||||
)
|
||||
model = "gemini-2.0-flash"
|
||||
|
||||
# Render the PDF page as an image (render_pdf_to_base64png is 1-indexed)
|
||||
try:
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=2048)
|
||||
except Exception as e:
|
||||
print(f"Error rendering PDF page: {str(e)}")
|
||||
return None
|
||||
|
||||
image_part = types.Part(inline_data=types.Blob(mime_type="image/png", data=base64.b64decode(image_base64)))
|
||||
|
||||
# Prepare prompt for Gemini to analyze document layout
|
||||
contents = [
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
image_part,
|
||||
types.Part.from_text(
|
||||
text=(
|
||||
"Please answer the following questions about the document in JSON format:\n"
|
||||
"-How many columns are used in the main text document layout?\n"
|
||||
"-How many unique articles are captured in main text on this page?\n"
|
||||
"-Are there any text inserts in the main article content?\n"
|
||||
"-Do any of the main content articles start with a dropcap?\n"
|
||||
"-Are there any boxed out regions of text that need to be read separately from the main article content?\n"
|
||||
"-Are there any regions of text with a different orientation/rotation?"
|
||||
)
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
generate_content_config = types.GenerateContentConfig(
|
||||
temperature=0.2,
|
||||
top_p=0.95,
|
||||
top_k=40,
|
||||
max_output_tokens=2048,
|
||||
response_mime_type="application/json",
|
||||
response_schema=types.Schema(
|
||||
type=types.Type.OBJECT,
|
||||
required=[
|
||||
"num_columns",
|
||||
"num_unique_articles",
|
||||
"contains_text_inserts",
|
||||
"contains_dropcaps",
|
||||
"contains_boxed_regions",
|
||||
"contains_text_different_orientation",
|
||||
],
|
||||
properties={
|
||||
"num_columns": types.Schema(
|
||||
type=types.Type.INTEGER,
|
||||
),
|
||||
"num_unique_articles": types.Schema(
|
||||
type=types.Type.INTEGER,
|
||||
),
|
||||
"contains_text_inserts": types.Schema(
|
||||
type=types.Type.BOOLEAN,
|
||||
),
|
||||
"contains_dropcaps": types.Schema(
|
||||
type=types.Type.BOOLEAN,
|
||||
),
|
||||
"contains_boxed_regions": types.Schema(
|
||||
type=types.Type.BOOLEAN,
|
||||
),
|
||||
"contains_text_different_orientation": types.Schema(
|
||||
type=types.Type.BOOLEAN,
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
# Call Gemini API
|
||||
response = client.models.generate_content(model=model, contents=contents, config=generate_content_config)
|
||||
|
||||
print(response)
|
||||
|
||||
if not response.candidates or len(response.candidates) == 0:
|
||||
print(f"No response generated for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
if response.candidates[0].finish_reason != types.FinishReason.STOP:
|
||||
print(f"Response generation incomplete for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
# Parse the response
|
||||
response_text = response.candidates[0].content.parts[0].text
|
||||
|
||||
layout_analysis = json.loads(response_text)
|
||||
|
||||
print(f"Layout analysis for {pdf_path} page {page_num}:")
|
||||
print(json.dumps(layout_analysis, indent=2))
|
||||
|
||||
# Return both the layout analysis and the rendered image (base64 string)
|
||||
return (layout_analysis, image_base64)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error analyzing document layout in {pdf_path} page {page_num}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def extract_document_content(pdf_path: str, page_num: int, image_base64: str, api_key: str) -> Optional[str]:
|
||||
"""
|
||||
Use Gemini to extract full document content from a rendered PDF page.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
image_base64: The base64 string of the rendered page image
|
||||
api_key: Gemini API key
|
||||
|
||||
Returns:
|
||||
Optional[str]: The extracted document content in markdown format, or None if extraction fails.
|
||||
"""
|
||||
# Initialize Gemini client
|
||||
client = genai.Client(
|
||||
api_key=api_key,
|
||||
)
|
||||
model = "gemini-2.0-flash"
|
||||
|
||||
image_part = types.Part(inline_data=types.Blob(mime_type="image/png", data=base64.b64decode(image_base64)))
|
||||
|
||||
# Prepare prompt for Gemini to extract document content
|
||||
contents = [
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
image_part,
|
||||
types.Part.from_text(
|
||||
text=(
|
||||
"Analyze the document attached and output it in markdown format. "
|
||||
"Output equations as Latex escaped with $$. "
|
||||
"Output tables in HTML format that preserves the structure and content exactly, do not use <br> tags. "
|
||||
"Instead of the markdown table format, be sure to output tables in HTML, even though the rest of the document is styled in markdown. "
|
||||
"Output figures with just a simple markdown image placeholder."
|
||||
)
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
generate_content_config = types.GenerateContentConfig(temperature=0.2, top_p=0.95, top_k=40, max_output_tokens=8192)
|
||||
|
||||
try:
|
||||
# Call Gemini API
|
||||
response = client.models.generate_content(model=model, contents=contents, config=generate_content_config)
|
||||
|
||||
if not response.candidates or len(response.candidates) == 0:
|
||||
print(f"No response generated for content extraction in {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
if response.candidates[0].finish_reason != types.FinishReason.STOP:
|
||||
print(f"Content extraction incomplete for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
# Get the extracted content
|
||||
content = response.candidates[0].content.parts[0].text
|
||||
return content
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error extracting document content from {pdf_path} page {page_num}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def should_extract_full_content(layout_analysis: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Determine if full content extraction is needed based on layout analysis results.
|
||||
|
||||
Args:
|
||||
layout_analysis: Dictionary containing layout analysis results
|
||||
|
||||
Returns:
|
||||
bool: True if any of the special layout features are detected, False otherwise
|
||||
"""
|
||||
# Check for special layout features that warrant full content extraction
|
||||
features_to_check = ["text_inserts", "dropcaps", "boxed_regions", "rotated_text"]
|
||||
|
||||
# Also check if there are multiple columns or articles
|
||||
try:
|
||||
columns = layout_analysis.get("columns", 0)
|
||||
if isinstance(columns, str):
|
||||
columns = int(columns) if columns.isdigit() else 0
|
||||
|
||||
articles = layout_analysis.get("articles", 0)
|
||||
if isinstance(articles, str):
|
||||
articles = int(articles) if articles.isdigit() else 0
|
||||
|
||||
if columns > 1 or articles > 1:
|
||||
return True
|
||||
except (ValueError, TypeError):
|
||||
# If we can't parse the values, assume we need to extract
|
||||
pass
|
||||
|
||||
# Check for any True values in the features
|
||||
for feature in features_to_check:
|
||||
value = layout_analysis.get(feature, False)
|
||||
if isinstance(value, str):
|
||||
if value.lower() in ["yes", "true", "1"]:
|
||||
return True
|
||||
elif value:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def process_pdf(s3_path: str, temp_dir: str, output_dir: str, api_key: str) -> Dict:
|
||||
"""
|
||||
Process a single PDF from S3.
|
||||
|
||||
Args:
|
||||
s3_path: S3 path to the PDF
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: Gemini API key
|
||||
|
||||
Returns:
|
||||
Dict: Results of processing the PDF
|
||||
"""
|
||||
# Create a thread-specific temp directory to avoid conflicts
|
||||
thread_id = threading.get_ident()
|
||||
thread_temp_dir = os.path.join(temp_dir, f"thread_{thread_id}")
|
||||
os.makedirs(thread_temp_dir, exist_ok=True)
|
||||
|
||||
# Extract filename from S3 path
|
||||
pdf_filename = os.path.basename(s3_path)
|
||||
local_pdf_path = os.path.join(thread_temp_dir, pdf_filename)
|
||||
|
||||
# Download PDF from S3
|
||||
if not download_pdf_from_s3(s3_path, local_pdf_path):
|
||||
return {"error": f"Failed to download {s3_path}"}
|
||||
|
||||
pdf_filter = PdfFilter()
|
||||
|
||||
if pdf_filter.filter_out_pdf(local_pdf_path):
|
||||
print(f"Filtering out {pdf_filename}")
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
return {"error": f"PDF {pdf_filename} filtered out"}
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(local_pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return {"error": f"PDF {pdf_filename} has no pages"}
|
||||
|
||||
all_pages = list(range(len(reader.pages)))
|
||||
random.shuffle(all_pages)
|
||||
|
||||
results = {"filename": pdf_filename, "s3_path": s3_path}
|
||||
|
||||
for page_num in all_pages:
|
||||
# Analyze document layout
|
||||
layout_result = analyze_document_layout(local_pdf_path, page_num, api_key)
|
||||
if not layout_result:
|
||||
print(f"Failed to analyze layout in {pdf_filename} page {page_num+1}")
|
||||
continue
|
||||
|
||||
layout_analysis, image_base64 = layout_result
|
||||
results["layout_analysis"] = layout_analysis
|
||||
|
||||
# Determine if we need to extract full content
|
||||
full_extraction_needed = should_extract_full_content(layout_analysis)
|
||||
results["full_extraction_needed"] = full_extraction_needed
|
||||
|
||||
# Extract full content if needed
|
||||
if full_extraction_needed:
|
||||
content = extract_document_content(local_pdf_path, page_num, image_base64, api_key)
|
||||
results["content"] = content if content else "Content extraction failed"
|
||||
|
||||
# Extract the page and save to output dir
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(output_dir, "pdfs", f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
with file_lock: # Use lock when writing to shared output directory
|
||||
extract_page_from_pdf(local_pdf_path, output_pdf_path, page_num)
|
||||
|
||||
# Save analysis results
|
||||
output_json_path = os.path.join(output_dir, "results", f"{pdf_basename}_pg{page_num+1}.json")
|
||||
with file_lock:
|
||||
os.makedirs(os.path.join(output_dir, "results"), exist_ok=True)
|
||||
with open(output_json_path, "w") as f:
|
||||
json.dump(results, f, indent=2)
|
||||
|
||||
print(f"Processed {pdf_filename} page {page_num+1}, analysis saved to {output_json_path}")
|
||||
|
||||
# Process only one page per PDF
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
return {"error": f"Error processing {pdf_filename}: {str(e)}"}
|
||||
finally:
|
||||
# Cleanup
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
|
||||
|
||||
def process_pdfs_parallel(s3_paths: List[str], temp_dir: str, output_dir: str, api_key: str, max_docs: int, num_workers: int):
|
||||
"""
|
||||
Process PDFs in parallel using a thread pool.
|
||||
|
||||
Args:
|
||||
s3_paths: List of S3 paths to PDFs
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: Gemini API key
|
||||
max_docs: Maximum number of documents to process
|
||||
num_workers: Number of parallel workers to use
|
||||
"""
|
||||
# Create output directory structure
|
||||
os.makedirs(os.path.join(output_dir, "pdfs"), exist_ok=True)
|
||||
os.makedirs(os.path.join(output_dir, "results"), exist_ok=True)
|
||||
|
||||
# Create a summary file
|
||||
summary_file = os.path.join(output_dir, "summary.jsonl")
|
||||
|
||||
# Track processed documents
|
||||
processed_count = 0
|
||||
|
||||
# Create a ThreadPoolExecutor
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
|
||||
# Submit tasks and track futures
|
||||
futures = {executor.submit(process_pdf, s3_path, temp_dir, output_dir, api_key): s3_path for s3_path in s3_paths}
|
||||
|
||||
# Process results as they complete
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
s3_path = futures[future]
|
||||
try:
|
||||
# Get the result from this worker
|
||||
result = future.result()
|
||||
|
||||
# Add to summary file
|
||||
with file_lock:
|
||||
with open(summary_file, "a") as f:
|
||||
f.write(json.dumps(result) + "\n")
|
||||
|
||||
# Increment counter if no error
|
||||
if "error" not in result:
|
||||
processed_count += 1
|
||||
print(f"Successfully processed {os.path.basename(s3_path)}, total: {processed_count}")
|
||||
|
||||
# Check if we've reached the maximum number of documents
|
||||
if processed_count >= max_docs:
|
||||
print(f"Reached maximum number of documents ({max_docs}), stopping")
|
||||
# Cancel any pending futures
|
||||
for f in futures:
|
||||
if not f.done():
|
||||
f.cancel()
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
print(f"Task for {os.path.basename(s3_path)} generated an exception: {e}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Analyze document layout and extract content from PDF documents")
|
||||
parser.add_argument("--input_list", required=True, help="Path to a file containing S3 paths to PDFs")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to store extracted pages and analysis results")
|
||||
parser.add_argument("--api_key", help="Gemini API key (if not provided, will use GEMINI_API_KEY environment variable)")
|
||||
parser.add_argument("--temp_dir", default="/tmp/analyze_documents", help="Directory for temporary files")
|
||||
parser.add_argument("--max_docs", type=int, default=100, help="Maximum number of documents to process")
|
||||
parser.add_argument("--parallel", type=int, default=1, help="Number of parallel threads to use")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("GEMINI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: Gemini API key not provided. Use --api_key or set GEMINI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
os.makedirs(args.temp_dir, exist_ok=True)
|
||||
os.makedirs(os.path.join(args.output_dir, "pdfs"), exist_ok=True)
|
||||
os.makedirs(os.path.join(args.output_dir, "results"), exist_ok=True)
|
||||
|
||||
# Reservoir sampling implementation
|
||||
s3_paths = []
|
||||
with open(args.input_list, "r") as f:
|
||||
for i, line in enumerate(tqdm(f)):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
if i < 100000:
|
||||
s3_paths.append(line)
|
||||
else:
|
||||
# Randomly replace elements with decreasing probability
|
||||
j = random.randint(0, i)
|
||||
if j < 100000:
|
||||
s3_paths[j] = line
|
||||
|
||||
print(f"Found {len(s3_paths)} PDF paths in input list")
|
||||
|
||||
# Determine number of workers to use
|
||||
num_workers = max(1, min(args.parallel, len(s3_paths)))
|
||||
print(f"Processing PDFs using {num_workers} parallel workers")
|
||||
|
||||
# Process PDFs in parallel
|
||||
process_pdfs_parallel(s3_paths, args.temp_dir, args.output_dir, api_key, args.max_docs, num_workers)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,491 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_tables.py - Extract tables from PDF documents and create table tests.
|
||||
|
||||
This script:
|
||||
1. Takes a file containing S3 paths to PDF documents as input
|
||||
2. For each PDF, extracts a random page and renders it to an image
|
||||
3. Uses Gemini to identify tables in the rendered image
|
||||
4. Extracts table content and creates table relationship tests by making a second Gemini request
|
||||
that now includes the page image alongside the prompt (e.g., "Given cell with {cell_value}, which cell is directly to the left of it?")
|
||||
5. Extracts the page from the PDF and saves it to an output folder
|
||||
|
||||
Usage:
|
||||
python mine_tables.py --input_list path/to/s3_paths.txt --output_dir path/to/output --api_key your_gemini_api_key [--parallel 4]
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import concurrent.futures
|
||||
import os
|
||||
import random
|
||||
import threading
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import boto3
|
||||
import numpy as np
|
||||
import pypdf
|
||||
from bs4 import BeautifulSoup
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.bench.tests import TableTest, save_tests
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
|
||||
# Create a thread-safe lock for writing to the output file
|
||||
file_lock = threading.Lock()
|
||||
tests_lock = threading.Lock()
|
||||
|
||||
|
||||
def download_pdf_from_s3(s3_path: str, local_path: str) -> bool:
|
||||
"""
|
||||
Download a PDF file from S3.
|
||||
|
||||
Args:
|
||||
s3_path: The S3 path (s3://bucket/path/to/file.pdf)
|
||||
local_path: The local path to save the file
|
||||
|
||||
Returns:
|
||||
bool: True if download was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Parse S3 path
|
||||
parts = s3_path.replace("s3://", "").split("/", 1)
|
||||
bucket = parts[0]
|
||||
key = parts[1]
|
||||
|
||||
# Create S3 client
|
||||
s3 = boto3.client("s3")
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
|
||||
# Download file
|
||||
s3.download_file(bucket, key, local_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {s3_path}: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def extract_page_from_pdf(input_path: str, output_path: str, page_num: int) -> bool:
|
||||
"""
|
||||
Extract a specific page from a PDF and save it as a new PDF.
|
||||
|
||||
Args:
|
||||
input_path: Path to the input PDF
|
||||
output_path: Path to save the extracted page
|
||||
page_num: The page number to extract (0-indexed)
|
||||
|
||||
Returns:
|
||||
bool: True if extraction was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Ensure output directory exists
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
# Read the input PDF
|
||||
reader = pypdf.PdfReader(input_path)
|
||||
|
||||
# Check if page number is valid
|
||||
if page_num >= len(reader.pages):
|
||||
print(f"Page number {page_num} out of range for {input_path} with {len(reader.pages)} pages")
|
||||
return False
|
||||
|
||||
# Create a new PDF with just the selected page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error extracting page {page_num} from {input_path}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def detect_tables(pdf_path: str, page_num: int, api_key: str) -> Optional[Tuple[List[np.ndarray], str]]:
|
||||
"""
|
||||
Use Gemini to detect tables in a rendered PDF page.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: Gemini API key
|
||||
|
||||
Returns:
|
||||
Optional[Tuple[List[np.ndarray], str]]:
|
||||
A tuple with a list of detected tables (as numpy arrays) and the base64 string of the rendered page image.
|
||||
Returns None if detection fails.
|
||||
"""
|
||||
# Initialize Gemini client
|
||||
client = genai.Client(
|
||||
api_key=api_key,
|
||||
)
|
||||
model = "gemini-2.0-flash"
|
||||
|
||||
# Render the PDF page as an image (render_pdf_to_base64png is 1-indexed)
|
||||
try:
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=2048)
|
||||
except Exception as e:
|
||||
print(f"Error rendering PDF page: {str(e)}")
|
||||
return None
|
||||
|
||||
image_part = types.Part(inline_data=types.Blob(mime_type="image/png", data=base64.b64decode(image_base64)))
|
||||
|
||||
# Prepare prompt for Gemini to extract tables
|
||||
contents = [
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
image_part,
|
||||
types.Part.from_text(
|
||||
text=(
|
||||
"Analyze the document attached and output it in markdown format. "
|
||||
"Output equations as Latex escaped with $$. "
|
||||
"Output tables in HTML format that preserves the structure and content exactly, do not use <br> tags. "
|
||||
"Instead of the markdown table format, be sure to output tables in HTML, even though the rest of the document is styled in markdown. "
|
||||
"Output figures with just a simple markdown image placeholder."
|
||||
)
|
||||
),
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
generate_content_config = types.GenerateContentConfig(temperature=0.2, top_p=0.95, top_k=40, max_output_tokens=8192)
|
||||
|
||||
try:
|
||||
# Call Gemini API
|
||||
response = client.models.generate_content(model=model, contents=contents, config=generate_content_config)
|
||||
|
||||
if not response.candidates or len(response.candidates) == 0:
|
||||
print(f"No response generated for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
if response.candidates[0].finish_reason != types.FinishReason.STOP:
|
||||
print(f"Response generation incomplete for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
# Parse the response
|
||||
response_text = response.candidates[0].content.parts[0].text
|
||||
|
||||
print(response_text)
|
||||
|
||||
# Parse tables from HTML
|
||||
parsed_tables = []
|
||||
soup = BeautifulSoup(response_text, "html.parser")
|
||||
tables = soup.find_all("table")
|
||||
|
||||
for table in tables:
|
||||
rows = table.find_all("tr")
|
||||
table_data = []
|
||||
for row in rows:
|
||||
cells = row.find_all(["th", "td"])
|
||||
row_data = [cell.get_text().strip() for cell in cells]
|
||||
table_data.append(row_data)
|
||||
# Ensure all rows have the same number of columns
|
||||
if table_data:
|
||||
max_cols = max(len(row) for row in table_data)
|
||||
padded_data = [row + [""] * (max_cols - len(row)) for row in table_data]
|
||||
table_array = np.array(padded_data)
|
||||
parsed_tables.append(table_array)
|
||||
|
||||
# Return both the parsed tables and the rendered image (base64 string)
|
||||
return (parsed_tables, image_base64) if parsed_tables else None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error detecting tables in {pdf_path} page {page_num}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def generate_table_tests(tables: List[np.ndarray], pdf_image: str, api_key: str, max_tests_per_table: int = 3) -> List[Dict]:
|
||||
"""
|
||||
Generate table tests from the detected tables by making a second Gemini request for each candidate cell.
|
||||
|
||||
For each candidate cell in a table, the function selects one valid relationship (e.g., "left", "up", "top_heading", etc.)
|
||||
and sends a prompt to Gemini including the page image. For example:
|
||||
"Given a cell in a table with value 'XYZ', please answer: which cell is directly to the left of it? Provide only the cell's text."
|
||||
|
||||
Args:
|
||||
tables: List of tables as numpy arrays
|
||||
pdf_image: Base64 string of the rendered page image
|
||||
api_key: Gemini API key to use for generating relationship tests
|
||||
max_tests_per_table: Maximum number of tests to generate per table
|
||||
|
||||
Returns:
|
||||
List of table test dictionaries
|
||||
"""
|
||||
tests = []
|
||||
# Initialize Gemini client for test queries
|
||||
client = genai.Client(api_key=api_key)
|
||||
model = "gemini-2.0-flash"
|
||||
config = types.GenerateContentConfig(temperature=0.2, top_p=0.95, top_k=40, max_output_tokens=100)
|
||||
|
||||
# Mapping for relationship prompts
|
||||
prompt_map = {
|
||||
"up": "which cell is directly above it?",
|
||||
"down": "which cell is directly below it?",
|
||||
"left": "which cell is directly to the left of it?",
|
||||
"right": "which cell is directly to the right of it?",
|
||||
"top_heading": "what is the top heading (the heading for the column at the top of the table) for this cell?",
|
||||
"left_heading": "what is the left heading (the heading for this row on the left part of the table) for this cell?",
|
||||
}
|
||||
|
||||
# Create an image part from the rendered pdf image
|
||||
image_part = types.Part(inline_data=types.Blob(mime_type="image/png", data=base64.b64decode(pdf_image)))
|
||||
|
||||
for table in tables:
|
||||
rows, cols = table.shape
|
||||
if table.size == 0 or rows < 2 or cols < 2:
|
||||
continue # Skip tables that are too small
|
||||
|
||||
# Try up to 3x max_tests_per_table candidate cells
|
||||
candidate_positions = []
|
||||
for row in range(rows):
|
||||
for col in range(cols):
|
||||
if not table[row, col].strip():
|
||||
continue
|
||||
if row > 0:
|
||||
candidate_positions.append((row, col, "up"))
|
||||
if row < rows - 1:
|
||||
candidate_positions.append((row, col, "down"))
|
||||
if col > 0:
|
||||
candidate_positions.append((row, col, "left"))
|
||||
if col < cols - 1:
|
||||
candidate_positions.append((row, col, "right"))
|
||||
if row > 0:
|
||||
candidate_positions.append((row, col, "top_heading"))
|
||||
if col > 0:
|
||||
candidate_positions.append((row, col, "left_heading"))
|
||||
|
||||
random.shuffle(candidate_positions)
|
||||
tests_for_this_table = 0
|
||||
|
||||
for row, col, relationship in candidate_positions:
|
||||
if tests_for_this_table >= max_tests_per_table:
|
||||
break
|
||||
|
||||
cell_value = table[row, col].strip()
|
||||
|
||||
prompt = (
|
||||
f"Given a cell in a table with value '{cell_value}', please answer: "
|
||||
f"{prompt_map[relationship]} Provide only the cell's text or output 'null' if there is not a matching cell."
|
||||
)
|
||||
|
||||
try:
|
||||
contents = [types.Content(role="user", parts=[image_part, types.Part.from_text(text=prompt)])]
|
||||
response = client.models.generate_content(model=model, contents=contents, config=config)
|
||||
if not response.candidates or len(response.candidates) == 0 or response.candidates[0].finish_reason != types.FinishReason.STOP:
|
||||
continue
|
||||
answer_text = response.candidates[0].content.parts[0].text.strip()
|
||||
if answer_text and "null" not in answer_text:
|
||||
test_data = {"cell": cell_value, relationship: answer_text}
|
||||
tests.append(test_data)
|
||||
tests_for_this_table += 1
|
||||
except Exception as e:
|
||||
print(f"Error querying Gemini for cell '{cell_value}' and relationship '{relationship}': {str(e)}")
|
||||
|
||||
return tests
|
||||
|
||||
|
||||
def process_pdf(s3_path: str, temp_dir: str, output_dir: str, api_key: str) -> List[TableTest]:
|
||||
"""
|
||||
Process a single PDF from S3.
|
||||
|
||||
Args:
|
||||
s3_path: S3 path to the PDF
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: Gemini API key
|
||||
|
||||
Returns:
|
||||
List[TableTest]: List of generated table tests
|
||||
"""
|
||||
# Create a thread-specific temp directory to avoid conflicts
|
||||
thread_id = threading.get_ident()
|
||||
thread_temp_dir = os.path.join(temp_dir, f"thread_{thread_id}")
|
||||
os.makedirs(thread_temp_dir, exist_ok=True)
|
||||
|
||||
# Extract filename from S3 path
|
||||
pdf_filename = os.path.basename(s3_path)
|
||||
local_pdf_path = os.path.join(thread_temp_dir, pdf_filename)
|
||||
|
||||
# Download PDF from S3
|
||||
if not download_pdf_from_s3(s3_path, local_pdf_path):
|
||||
return []
|
||||
|
||||
pdf_filter = PdfFilter()
|
||||
|
||||
if pdf_filter.filter_out_pdf(local_pdf_path):
|
||||
print(f"Filtering out {pdf_filename}")
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
return []
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(local_pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return []
|
||||
|
||||
all_pages = list(range(len(reader.pages)))
|
||||
random.shuffle(all_pages)
|
||||
|
||||
local_tests = []
|
||||
|
||||
for page_num in all_pages:
|
||||
# Detect tables and obtain the rendered image for this page
|
||||
result = detect_tables(local_pdf_path, page_num, api_key)
|
||||
if not result:
|
||||
print(f"No tables detected in {pdf_filename} page {page_num+1}")
|
||||
continue
|
||||
|
||||
tables, image_base64 = result
|
||||
|
||||
# Generate table tests using the new Gemini query approach with the page image
|
||||
table_tests_data = generate_table_tests(tables, image_base64, api_key, max_tests_per_table=5)
|
||||
|
||||
if not table_tests_data:
|
||||
print(f"Could not generate valid tests for tables in {pdf_filename} page {page_num+1}")
|
||||
continue
|
||||
|
||||
# Extract the page and save to output dir
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(output_dir, "pdfs", f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
with file_lock: # Use lock when writing to shared output directory
|
||||
extract_page_from_pdf(local_pdf_path, output_pdf_path, page_num)
|
||||
|
||||
# Create table tests
|
||||
for i, test_data in enumerate(table_tests_data):
|
||||
test_id = f"{pdf_basename}_pg{page_num+1}_table_{i:02d}"
|
||||
test = TableTest(
|
||||
id=test_id,
|
||||
pdf=f"{pdf_basename}_pg{page_num+1}.pdf",
|
||||
page=1, # The extracted PDF has only one page
|
||||
type="table",
|
||||
cell=test_data["cell"],
|
||||
url=s3_path, # Added the S3 path as the url field
|
||||
up=test_data.get("up", None),
|
||||
down=test_data.get("down", None),
|
||||
left=test_data.get("left", None),
|
||||
right=test_data.get("right", None),
|
||||
top_heading=test_data.get("top_heading", None),
|
||||
left_heading=test_data.get("left_heading", None),
|
||||
)
|
||||
local_tests.append(test)
|
||||
|
||||
print(f"Processed {pdf_filename} page {page_num+1}, found {len(tables)} tables, created {len(table_tests_data)} tests")
|
||||
break # Process only one page per PDF
|
||||
|
||||
return local_tests
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
return []
|
||||
finally:
|
||||
# Cleanup
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
|
||||
|
||||
def process_pdfs_parallel(s3_paths: List[str], temp_dir: str, output_dir: str, api_key: str, max_tests: int, num_workers: int):
|
||||
"""
|
||||
Process PDFs in parallel using a thread pool.
|
||||
|
||||
Args:
|
||||
s3_paths: List of S3 paths to PDFs
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: Gemini API key
|
||||
max_tests: Maximum number of tests to generate
|
||||
num_workers: Number of parallel workers to use
|
||||
"""
|
||||
# Create shared resources
|
||||
all_tests = []
|
||||
output_file = os.path.join(output_dir, "table_tests.jsonl")
|
||||
|
||||
# Create a ThreadPoolExecutor
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
|
||||
# Submit tasks and track futures
|
||||
futures = {executor.submit(process_pdf, s3_path, temp_dir, output_dir, api_key): s3_path for s3_path in s3_paths}
|
||||
|
||||
# Process results as they complete
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
s3_path = futures[future]
|
||||
try:
|
||||
# Get the tests produced by this worker
|
||||
new_tests = future.result()
|
||||
|
||||
# If we got new tests, add them to our collection
|
||||
if new_tests:
|
||||
all_tests.extend(new_tests)
|
||||
save_tests(all_tests, output_file)
|
||||
print(f"Added {len(new_tests)} tests from {os.path.basename(s3_path)}, total: {len(all_tests)}")
|
||||
|
||||
# Check if we've reached the maximum number of tests
|
||||
if len(all_tests) >= max_tests:
|
||||
print(f"Reached maximum number of tests ({max_tests}), stopping")
|
||||
# Cancel any pending futures
|
||||
for f in futures:
|
||||
if not f.done():
|
||||
f.cancel()
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
print(f"Task for {os.path.basename(s3_path)} generated an exception: {e}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Extract tables from PDF documents and create table tests")
|
||||
parser.add_argument("--input_list", required=True, help="Path to a file containing S3 paths to PDFs")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to store extracted pages and tests")
|
||||
parser.add_argument("--api_key", help="Gemini API key (if not provided, will use GEMINI_API_KEY environment variable)")
|
||||
parser.add_argument("--temp_dir", default="/tmp/mine_tables", help="Directory for temporary files")
|
||||
parser.add_argument("--max_tests", type=int, default=100, help="Maximum number of tests to generate")
|
||||
parser.add_argument("--parallel", type=int, default=1, help="Number of parallel threads to use")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("GEMINI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: Gemini API key not provided. Use --api_key or set GEMINI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
os.makedirs(args.temp_dir, exist_ok=True)
|
||||
os.makedirs(os.path.join(args.output_dir, "pdfs"), exist_ok=True)
|
||||
|
||||
# Reservoir sampling implementation
|
||||
s3_paths = []
|
||||
with open(args.input_list, "r") as f:
|
||||
for i, line in enumerate(tqdm(f)):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
if i < 100000:
|
||||
s3_paths.append(line)
|
||||
else:
|
||||
# Randomly replace elements with decreasing probability
|
||||
j = random.randint(0, i)
|
||||
if j < 100000:
|
||||
s3_paths[j] = line
|
||||
|
||||
print(f"Found {len(s3_paths)} PDF paths in input list")
|
||||
|
||||
# Determine number of workers to use
|
||||
num_workers = max(1, min(args.parallel, len(s3_paths)))
|
||||
print(f"Processing PDFs using {num_workers} parallel workers")
|
||||
|
||||
# Process PDFs in parallel
|
||||
process_pdfs_parallel(s3_paths, args.temp_dir, args.output_dir, api_key, args.max_tests, num_workers)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,418 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_tables.py - Extract tables from PDF documents and create table tests.
|
||||
|
||||
This script:
|
||||
1. Takes a file containing S3 paths to PDF documents as input
|
||||
2. For each PDF, extracts a random page and renders it to an image
|
||||
3. Uses GPT-4o to identify tables in the rendered image
|
||||
4. Extracts table content and creates table relationship tests by making a second GPT-4o request
|
||||
that now includes the page image alongside the prompt (e.g., "Given cell with {cell_value}, which cell is directly to the left of it?")
|
||||
5. Extracts the page from the PDF and saves it to an output folder
|
||||
|
||||
Usage:
|
||||
python mine_tables.py --input_list path/to/s3_paths.txt --output_dir path/to/output --api_key your_openai_api_key
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import random
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import boto3
|
||||
import numpy as np
|
||||
import pypdf
|
||||
from bs4 import BeautifulSoup
|
||||
from openai import OpenAI
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.bench.tests import TableTest, save_tests
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
|
||||
|
||||
def download_pdf_from_s3(s3_path: str, local_path: str) -> bool:
|
||||
"""
|
||||
Download a PDF file from S3.
|
||||
|
||||
Args:
|
||||
s3_path: The S3 path (s3://bucket/path/to/file.pdf)
|
||||
local_path: The local path to save the file
|
||||
|
||||
Returns:
|
||||
bool: True if download was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Parse S3 path
|
||||
parts = s3_path.replace("s3://", "").split("/", 1)
|
||||
bucket = parts[0]
|
||||
key = parts[1]
|
||||
|
||||
# Create S3 client
|
||||
s3 = boto3.client("s3")
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
|
||||
# Download file
|
||||
s3.download_file(bucket, key, local_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {s3_path}: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def extract_page_from_pdf(input_path: str, output_path: str, page_num: int) -> bool:
|
||||
"""
|
||||
Extract a specific page from a PDF and save it as a new PDF.
|
||||
|
||||
Args:
|
||||
input_path: Path to the input PDF
|
||||
output_path: Path to save the extracted page
|
||||
page_num: The page number to extract (0-indexed)
|
||||
|
||||
Returns:
|
||||
bool: True if extraction was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Ensure output directory exists
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
# Read the input PDF
|
||||
reader = pypdf.PdfReader(input_path)
|
||||
|
||||
# Check if page number is valid
|
||||
if page_num >= len(reader.pages):
|
||||
print(f"Page number {page_num} out of range for {input_path} with {len(reader.pages)} pages")
|
||||
return False
|
||||
|
||||
# Create a new PDF with just the selected page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error extracting page {page_num} from {input_path}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def detect_tables(pdf_path: str, page_num: int, api_key: str) -> Optional[Tuple[List[np.ndarray], str]]:
|
||||
"""
|
||||
Use GPT-4o to detect tables in a rendered PDF page.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
Optional[Tuple[List[np.ndarray], str]]:
|
||||
A tuple with a list of detected tables (as numpy arrays) and the base64 string of the rendered page image.
|
||||
Returns None if detection fails.
|
||||
"""
|
||||
# Initialize OpenAI client
|
||||
client = OpenAI(api_key=api_key)
|
||||
model = "gpt-4o"
|
||||
|
||||
# Render the PDF page as an image (render_pdf_to_base64png is 1-indexed)
|
||||
try:
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=2048)
|
||||
except Exception as e:
|
||||
print(f"Error rendering PDF page: {str(e)}")
|
||||
return None
|
||||
|
||||
# Prepare prompt for GPT-4o to extract tables
|
||||
try:
|
||||
# Call OpenAI API
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}", "detail": "high"}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"Analyze the document attached and output it in markdown format. "
|
||||
"Output equations as Latex escaped with $$. "
|
||||
"Output tables in valid HTML format that preserves the structure and content exactly. "
|
||||
"Output figures with just a simple markdown image placeholder."
|
||||
),
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
temperature=0.2,
|
||||
)
|
||||
|
||||
if not response.choices or len(response.choices) == 0:
|
||||
print(f"No response generated for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
# Parse the response
|
||||
response_text = response.choices[0].message.content
|
||||
|
||||
print(response_text)
|
||||
|
||||
# Parse tables from HTML
|
||||
parsed_tables = []
|
||||
soup = BeautifulSoup(response_text, "html.parser")
|
||||
tables = soup.find_all("table")
|
||||
|
||||
for table in tables:
|
||||
rows = table.find_all("tr")
|
||||
table_data = []
|
||||
for row in rows:
|
||||
cells = row.find_all(["th", "td"])
|
||||
row_data = [cell.get_text().strip() for cell in cells]
|
||||
table_data.append(row_data)
|
||||
# Ensure all rows have the same number of columns
|
||||
if table_data:
|
||||
max_cols = max(len(row) for row in table_data)
|
||||
padded_data = [row + [""] * (max_cols - len(row)) for row in table_data]
|
||||
table_array = np.array(padded_data)
|
||||
parsed_tables.append(table_array)
|
||||
|
||||
# Return both the parsed tables and the rendered image (base64 string)
|
||||
return (parsed_tables, image_base64) if parsed_tables else None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error detecting tables in {pdf_path} page {page_num}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def generate_table_tests(tables: List[np.ndarray], pdf_image: str, api_key: str, max_tests_per_table: int = 3) -> List[Dict]:
|
||||
"""
|
||||
Generate table tests from the detected tables by making a second GPT-4o request for each candidate cell.
|
||||
|
||||
For each candidate cell in a table, the function selects one valid relationship (e.g., "left", "up", "top_heading", etc.)
|
||||
and sends a prompt to GPT-4o including the page image. For example:
|
||||
"Given a cell in a table with value 'XYZ', please answer: which cell is directly to the left of it? Provide only the cell's text."
|
||||
|
||||
Args:
|
||||
tables: List of tables as numpy arrays
|
||||
pdf_image: Base64 string of the rendered page image
|
||||
api_key: OpenAI API key to use for generating relationship tests
|
||||
max_tests_per_table: Maximum number of tests to generate per table
|
||||
|
||||
Returns:
|
||||
List of table test dictionaries
|
||||
"""
|
||||
tests = []
|
||||
# Initialize OpenAI client for test queries
|
||||
client = OpenAI(api_key=api_key)
|
||||
model = "gpt-4o"
|
||||
|
||||
# Mapping for relationship prompts
|
||||
prompt_map = {
|
||||
"up": "which cell is directly above it?",
|
||||
"down": "which cell is directly below it?",
|
||||
"left": "which cell is directly to the left of it?",
|
||||
"right": "which cell is directly to the right of it?",
|
||||
"top_heading": "what is the top heading for this cell?",
|
||||
"left_heading": "what is the left heading for this cell?",
|
||||
}
|
||||
|
||||
for table in tables:
|
||||
rows, cols = table.shape
|
||||
if table.size == 0 or rows < 2 or cols < 2:
|
||||
continue # Skip tables that are too small
|
||||
|
||||
# Try up to 3x max_tests_per_table candidate cells
|
||||
candidate_positions = []
|
||||
for _ in range(max_tests_per_table * 3):
|
||||
row = random.randint(0, rows - 1)
|
||||
col = random.randint(0, cols - 1)
|
||||
if not table[row, col].strip():
|
||||
continue
|
||||
candidate_positions.append((row, col))
|
||||
|
||||
random.shuffle(candidate_positions)
|
||||
tests_for_this_table = 0
|
||||
|
||||
for row, col in candidate_positions:
|
||||
if tests_for_this_table >= max_tests_per_table:
|
||||
break
|
||||
|
||||
cell_value = table[row, col].strip()
|
||||
# Determine valid relationship types based on candidate's position
|
||||
valid_relationships = []
|
||||
if row > 0:
|
||||
valid_relationships.append("up")
|
||||
if row < rows - 1:
|
||||
valid_relationships.append("down")
|
||||
if col > 0:
|
||||
valid_relationships.append("left")
|
||||
if col < cols - 1:
|
||||
valid_relationships.append("right")
|
||||
if row > 0:
|
||||
valid_relationships.append("top_heading")
|
||||
if col > 0:
|
||||
valid_relationships.append("left_heading")
|
||||
if not valid_relationships:
|
||||
continue
|
||||
|
||||
relationship = random.choice(valid_relationships)
|
||||
prompt = (
|
||||
f"Given a cell in a table with value '{cell_value}', please answer: "
|
||||
f"{prompt_map[relationship]} Provide only the cell's text or output 'null' if there is not a matching cell."
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{pdf_image}", "detail": "high"}},
|
||||
{"type": "text", "text": prompt},
|
||||
],
|
||||
}
|
||||
],
|
||||
temperature=0.2,
|
||||
)
|
||||
|
||||
if not response.choices or len(response.choices) == 0:
|
||||
continue
|
||||
|
||||
answer_text = response.choices[0].message.content.strip()
|
||||
if answer_text and "null" not in answer_text:
|
||||
test_data = {"cell": cell_value, relationship: answer_text}
|
||||
tests.append(test_data)
|
||||
tests_for_this_table += 1
|
||||
except Exception as e:
|
||||
print(f"Error querying GPT-4o for cell '{cell_value}' and relationship '{relationship}': {str(e)}")
|
||||
|
||||
return tests
|
||||
|
||||
|
||||
def process_pdf(s3_path: str, temp_dir: str, output_dir: str, api_key: str, tests: List[TableTest]) -> None:
|
||||
"""
|
||||
Process a single PDF from S3.
|
||||
|
||||
Args:
|
||||
s3_path: S3 path to the PDF
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: OpenAI API key
|
||||
tests: List to append tests to
|
||||
"""
|
||||
# Extract filename from S3 path
|
||||
pdf_filename = os.path.basename(s3_path)
|
||||
local_pdf_path = os.path.join(temp_dir, pdf_filename)
|
||||
|
||||
# Download PDF from S3
|
||||
if not download_pdf_from_s3(s3_path, local_pdf_path):
|
||||
return
|
||||
|
||||
pdf_filter = PdfFilter()
|
||||
|
||||
if pdf_filter.filter_out_pdf(local_pdf_path):
|
||||
print(f"Filtering out {pdf_filename}")
|
||||
return
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(local_pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return
|
||||
|
||||
all_pages = list(range(len(reader.pages)))
|
||||
random.shuffle(all_pages)
|
||||
|
||||
for page_num in all_pages:
|
||||
# Detect tables and obtain the rendered image for this page
|
||||
result = detect_tables(local_pdf_path, page_num, api_key)
|
||||
if not result:
|
||||
print(f"No tables detected in {pdf_filename} page {page_num+1}")
|
||||
continue
|
||||
|
||||
tables, image_base64 = result
|
||||
|
||||
# Generate table tests using the new GPT-4o query approach with the page image
|
||||
table_tests_data = generate_table_tests(tables, image_base64, api_key, max_tests_per_table=5)
|
||||
|
||||
if not table_tests_data:
|
||||
print(f"Could not generate valid tests for tables in {pdf_filename} page {page_num+1}")
|
||||
continue
|
||||
|
||||
# Extract the page and save to output dir
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(output_dir, "pdfs", f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
extract_page_from_pdf(local_pdf_path, output_pdf_path, page_num)
|
||||
|
||||
# Create table tests
|
||||
for i, test_data in enumerate(table_tests_data):
|
||||
test_id = f"{pdf_basename}_pg{page_num+1}_table_{i:02d}"
|
||||
test = TableTest(
|
||||
id=test_id,
|
||||
pdf=f"{pdf_basename}_pg{page_num+1}.pdf",
|
||||
page=1, # The extracted PDF has only one page
|
||||
type="table",
|
||||
cell=test_data["cell"],
|
||||
up=test_data.get("up", None),
|
||||
down=test_data.get("down", None),
|
||||
left=test_data.get("left", None),
|
||||
right=test_data.get("right", None),
|
||||
top_heading=test_data.get("top_heading", None),
|
||||
left_heading=test_data.get("left_heading", None),
|
||||
)
|
||||
tests.append(test)
|
||||
|
||||
print(f"Processed {pdf_filename} page {page_num+1}, found {len(tables)} tables, created {len(table_tests_data)} tests")
|
||||
return # Process only one page per PDF
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
finally:
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Extract tables from PDF documents and create table tests")
|
||||
parser.add_argument("--input_list", required=True, help="Path to a file containing S3 paths to PDFs")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to store extracted pages and tests")
|
||||
parser.add_argument("--api_key", help="OpenAI API key (if not provided, will use OPENAI_API_KEY environment variable)")
|
||||
parser.add_argument("--temp_dir", default="/tmp/mine_tables", help="Directory for temporary files")
|
||||
parser.add_argument("--max_tests", type=int, default=100, help="Maximum number of tests to generate")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: OpenAI API key not provided. Use --api_key or set OPENAI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
os.makedirs(args.temp_dir, exist_ok=True)
|
||||
os.makedirs(os.path.join(args.output_dir, "pdfs"), exist_ok=True)
|
||||
|
||||
with open(args.input_list, "r") as f:
|
||||
s3_paths = [line.strip() for line in f if line.strip()]
|
||||
|
||||
print(f"Found {len(s3_paths)} PDF paths in input list")
|
||||
tests = []
|
||||
for s3_path in tqdm(s3_paths, desc="Processing PDFs"):
|
||||
process_pdf(s3_path, args.temp_dir, args.output_dir, api_key, tests)
|
||||
|
||||
if tests:
|
||||
save_tests(tests, os.path.join(args.output_dir, "table_tests.jsonl"))
|
||||
|
||||
if len(tests) >= args.max_tests:
|
||||
print(f"Reached maximum number of tests ({args.max_tests}), stopping")
|
||||
break
|
||||
|
||||
print(f"Saved {len(tests)} table tests to {os.path.join(args.output_dir, 'table_tests.jsonl')}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,333 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
mine_tables_gpt_simple.py - Identify PDF documents with tables and copy them.
|
||||
|
||||
This script:
|
||||
1. Takes a file containing S3 paths to PDF documents as input
|
||||
2. For each PDF, renders a random page and uses GPT-4o to check for tables
|
||||
3. Identifies PDFs where the page contains a table
|
||||
4. Copies those PDF files to a new output folder
|
||||
|
||||
Usage:
|
||||
python mine_tables_gpt_simple.py --input_list path/to/s3_paths.txt --output_dir path/to/output --api_key your_openai_api_key
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import random
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import Optional
|
||||
|
||||
import boto3
|
||||
import pypdf
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.filter import PdfFilter
|
||||
|
||||
TARGET_IMAGE_DIM = 1024
|
||||
|
||||
|
||||
class TableInfo(BaseModel):
|
||||
"""Information about a single table."""
|
||||
|
||||
num_rows: int
|
||||
num_cols: int
|
||||
|
||||
|
||||
class TableDetectionResponse(BaseModel):
|
||||
"""Structured output for table detection."""
|
||||
|
||||
tables: list[TableInfo]
|
||||
|
||||
|
||||
def download_pdf_from_s3(s3_path: str, local_path: str) -> bool:
|
||||
"""
|
||||
Download a PDF file from S3.
|
||||
|
||||
Args:
|
||||
s3_path: The S3 path (s3://bucket/path/to/file.pdf)
|
||||
local_path: The local path to save the file
|
||||
|
||||
Returns:
|
||||
bool: True if download was successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Parse S3 path
|
||||
parts = s3_path.replace("s3://", "").split("/", 1)
|
||||
bucket = parts[0]
|
||||
key = parts[1]
|
||||
|
||||
# Create S3 client
|
||||
s3 = boto3.client("s3")
|
||||
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
|
||||
# Download file
|
||||
s3.download_file(bucket, key, local_path)
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {s3_path}: {str(e)}")
|
||||
return False
|
||||
|
||||
|
||||
def get_cell_count_bucket(total_cells: int) -> str:
|
||||
"""
|
||||
Get the folder name for a given cell count, bucketed by powers of 2.
|
||||
|
||||
Args:
|
||||
total_cells: Total number of cells across all tables
|
||||
|
||||
Returns:
|
||||
str: Folder name like "0_cells", "1_cell", "2_cells", "4_cells", etc.
|
||||
"""
|
||||
if total_cells == 0:
|
||||
return "0_cells"
|
||||
elif total_cells == 1:
|
||||
return "1_cell"
|
||||
else:
|
||||
# Find the next power of 2 >= total_cells
|
||||
power = 1
|
||||
while power < total_cells:
|
||||
power *= 2
|
||||
return f"{power}_cells"
|
||||
|
||||
|
||||
def check_for_table(pdf_path: str, page_num: int, api_key: str) -> Optional[tuple[bool, int]]:
|
||||
"""
|
||||
Use GPT-4o to check if a page contains a table.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: The page number to analyze (0-indexed)
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
Optional[tuple[bool, int]]: Tuple of (has_table, total_cells) or None if detection fails
|
||||
"""
|
||||
# Initialize OpenAI client
|
||||
client = OpenAI(api_key=api_key)
|
||||
|
||||
try:
|
||||
# Render the PDF page as an image (render_pdf_to_base64png is 1-indexed)
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num + 1, target_longest_image_dim=TARGET_IMAGE_DIM)
|
||||
|
||||
# Prompt asking for detailed table information
|
||||
prompt = "Identify all tables on this page. For each table, count the number of rows and columns. Return an empty list if there are no tables."
|
||||
|
||||
response = client.beta.chat.completions.parse(
|
||||
model="gpt-5.1",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}}],
|
||||
}
|
||||
],
|
||||
max_completion_tokens=1000,
|
||||
response_format=TableDetectionResponse,
|
||||
)
|
||||
|
||||
if not response.choices or len(response.choices) == 0:
|
||||
print(f"No response generated for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
# Parse the structured response
|
||||
parsed_response = response.choices[0].message.parsed
|
||||
|
||||
if parsed_response is None:
|
||||
print(f"Failed to parse response for {pdf_path} page {page_num}")
|
||||
return None
|
||||
|
||||
tables = parsed_response.tables
|
||||
has_table = len(tables) > 0
|
||||
total_cells = sum(table.num_rows * table.num_cols for table in tables)
|
||||
|
||||
if has_table:
|
||||
print(f"Found {len(tables)} table(s) in {pdf_path} page {page_num + 1}, total cells: {total_cells}")
|
||||
for i, table in enumerate(tables, 1):
|
||||
print(f" Table {i}: {table.num_rows} rows × {table.num_cols} cols = {table.num_rows * table.num_cols} cells")
|
||||
|
||||
return (has_table, total_cells)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error checking {pdf_path} page {page_num}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
def process_pdf(s3_path: str, temp_dir: str, output_dir: str, api_key: str) -> bool:
|
||||
"""
|
||||
Process a single PDF from S3.
|
||||
|
||||
Args:
|
||||
s3_path: S3 path to the PDF
|
||||
temp_dir: Directory for temporary files
|
||||
output_dir: Directory for output files
|
||||
api_key: OpenAI API key
|
||||
|
||||
Returns:
|
||||
bool: True if the PDF has a table and was copied, False otherwise
|
||||
"""
|
||||
# Extract filename from S3 path
|
||||
pdf_filename = os.path.basename(s3_path)
|
||||
local_pdf_path = os.path.join(temp_dir, pdf_filename)
|
||||
|
||||
# Download PDF from S3
|
||||
if not download_pdf_from_s3(s3_path, local_pdf_path):
|
||||
return False
|
||||
|
||||
pdf_filter = PdfFilter()
|
||||
|
||||
if pdf_filter.filter_out_pdf(local_pdf_path):
|
||||
print(f"Filtering out {pdf_filename}")
|
||||
return False
|
||||
|
||||
try:
|
||||
# Read the PDF to get the number of pages
|
||||
reader = pypdf.PdfReader(local_pdf_path)
|
||||
num_pages = len(reader.pages)
|
||||
|
||||
if num_pages == 0:
|
||||
print(f"PDF {pdf_filename} has no pages")
|
||||
return False
|
||||
|
||||
# Select a random page to check
|
||||
page_num = random.randint(0, num_pages - 1)
|
||||
page_num = random.choice([page_num, 0]) # Bias 50% of the time to do the first page
|
||||
|
||||
# Check if the page contains a table
|
||||
result = check_for_table(local_pdf_path, page_num, api_key)
|
||||
|
||||
if result is None:
|
||||
return False
|
||||
|
||||
has_table, total_cells = result
|
||||
|
||||
if has_table:
|
||||
# Get the cell count bucket for organizing output
|
||||
bucket_name = get_cell_count_bucket(total_cells)
|
||||
bucket_dir = os.path.join(output_dir, bucket_name)
|
||||
os.makedirs(bucket_dir, exist_ok=True)
|
||||
|
||||
# Create output filename with basename_pgnum.pdf format
|
||||
pdf_basename = os.path.splitext(pdf_filename)[0]
|
||||
output_pdf_path = os.path.join(bucket_dir, f"{pdf_basename}_pg{page_num+1}.pdf")
|
||||
|
||||
# Extract the single page
|
||||
writer = pypdf.PdfWriter()
|
||||
writer.add_page(reader.pages[page_num])
|
||||
|
||||
# Write the output PDF
|
||||
with open(output_pdf_path, "wb") as output_file:
|
||||
writer.write(output_file)
|
||||
|
||||
print(f"Extracted page {page_num+1} with table from {pdf_filename} to {bucket_name}/{os.path.basename(output_pdf_path)}")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {pdf_filename}: {str(e)}")
|
||||
return False
|
||||
finally:
|
||||
if os.path.exists(local_pdf_path):
|
||||
os.remove(local_pdf_path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Identify and copy PDFs with tables")
|
||||
parser.add_argument("--input_list", required=True, help="Path to a file containing S3 paths to PDFs")
|
||||
parser.add_argument("--output_dir", required=True, help="Directory to copy PDFs with tables")
|
||||
parser.add_argument("--api_key", help="OpenAI API key (if not provided, will use OPENAI_API_KEY environment variable)")
|
||||
parser.add_argument("--temp_dir", default="/tmp/mine_tables", help="Directory for temporary files")
|
||||
parser.add_argument("--max_pdfs", type=int, default=100, help="Maximum number of PDFs with tables to find")
|
||||
parser.add_argument("--parallel", type=int, default=1, help="Number of parallel workers (default: 1 for sequential)")
|
||||
parser.add_argument("--reservoir_multiplier", type=int, default=100, help="Multiplier for reservoir sampling (default: 100x max_pdfs)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get API key
|
||||
api_key = args.api_key or os.environ.get("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
print("Error: OpenAI API key not provided. Use --api_key or set OPENAI_API_KEY environment variable.")
|
||||
return
|
||||
|
||||
os.makedirs(args.temp_dir, exist_ok=True)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Reservoir sampling to get random subset of PDFs
|
||||
reservoir_size = args.max_pdfs * args.reservoir_multiplier
|
||||
pdf_paths = []
|
||||
n = 0 # Total number of items seen
|
||||
|
||||
print(f"Using reservoir sampling with size {reservoir_size}")
|
||||
|
||||
with open(args.input_list, "r") as f:
|
||||
for line in tqdm(f):
|
||||
n += 1
|
||||
path = line.strip()
|
||||
if not path:
|
||||
continue
|
||||
|
||||
if len(pdf_paths) < reservoir_size:
|
||||
pdf_paths.append(path)
|
||||
else:
|
||||
# Randomly decide whether to include this item
|
||||
s = random.randint(1, n)
|
||||
if s <= reservoir_size:
|
||||
pdf_paths[s - 1] = path
|
||||
|
||||
# Shuffle the reservoir
|
||||
random.shuffle(pdf_paths)
|
||||
|
||||
print(f"Sampled {len(pdf_paths)} PDF paths from {n} total paths")
|
||||
|
||||
table_pdfs_found = 0
|
||||
|
||||
if args.parallel > 1:
|
||||
# Parallel processing
|
||||
print(f"Processing PDFs with {args.parallel} parallel workers")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=args.parallel) as executor:
|
||||
futures = []
|
||||
|
||||
# Submit all tasks
|
||||
for s3_path in pdf_paths:
|
||||
if table_pdfs_found >= args.max_pdfs:
|
||||
break
|
||||
future = executor.submit(process_pdf, s3_path, args.temp_dir, args.output_dir, api_key)
|
||||
futures.append(future)
|
||||
|
||||
# Process results as they complete
|
||||
with tqdm(total=min(len(pdf_paths), args.max_pdfs), desc="Processing PDFs") as pbar:
|
||||
for future in as_completed(futures):
|
||||
try:
|
||||
result = future.result()
|
||||
if result:
|
||||
table_pdfs_found += 1
|
||||
pbar.update(1)
|
||||
|
||||
if table_pdfs_found >= args.max_pdfs:
|
||||
print(f"Reached maximum number of PDFs with tables ({args.max_pdfs}), stopping")
|
||||
# Cancel remaining futures
|
||||
for f in futures:
|
||||
f.cancel()
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"Error in parallel processing: {str(e)}")
|
||||
else:
|
||||
# Sequential processing
|
||||
for s3_path in tqdm(pdf_paths, desc="Processing PDFs"):
|
||||
if process_pdf(s3_path, args.temp_dir, args.output_dir, api_key):
|
||||
table_pdfs_found += 1
|
||||
|
||||
if table_pdfs_found >= args.max_pdfs:
|
||||
print(f"Reached maximum number of PDFs with tables ({args.max_pdfs}), stopping")
|
||||
break
|
||||
|
||||
print(f"Found and copied {table_pdfs_found} PDFs with tables to {args.output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+162
@@ -0,0 +1,162 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
pick_mediod.py - Identify representative examples from repeated OCR outputs
|
||||
|
||||
This code will take as arguments two directories:
|
||||
--input and --output
|
||||
Each of those is going to be a directory that was generated by convert.py and is a candidate to be evaluated as part of benchmark.py
|
||||
What it will do is find and group all of the .md files into their repeats
|
||||
ex. input_dir/tables/buildingnotes_pg1_repeat1.md, input_dir/tables/buildingnotes_pg1_repeat2.md, etc.
|
||||
Then, for each repeat, it will use string similarity metrics to calculate the edit distance to every other repeat
|
||||
The repeat with the lowest mean edit distance will then get output as ..._repeat1.md in the output folder
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
from typing import Dict, List
|
||||
|
||||
from rapidfuzz import distance as fuzz_distance
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def compute_distance(text1: str, text2: str) -> float:
|
||||
"""
|
||||
Compute the edit distance between two text strings using rapidfuzz.
|
||||
Returns a normalized distance between 0.0 (identical) and 1.0 (completely different).
|
||||
"""
|
||||
# Use Levenshtein distance for string comparison
|
||||
return fuzz_distance.Levenshtein.normalized_distance(text1, text2)
|
||||
|
||||
|
||||
def find_mediod(texts: List[str]) -> int:
|
||||
"""
|
||||
Find the index of the mediod from a list of texts.
|
||||
The mediod is the text with the minimum average distance to all other texts.
|
||||
"""
|
||||
if not texts:
|
||||
return -1
|
||||
|
||||
if len(texts) == 1:
|
||||
return 0
|
||||
|
||||
# Calculate pairwise distances between all texts
|
||||
n = len(texts)
|
||||
distances = [[0.0 for _ in range(n)] for _ in range(n)]
|
||||
|
||||
for i in range(n):
|
||||
for j in range(i + 1, n):
|
||||
dist = compute_distance(texts[i], texts[j])
|
||||
distances[i][j] = dist
|
||||
distances[j][i] = dist
|
||||
|
||||
# Calculate average distance of each text to all others
|
||||
avg_distances = []
|
||||
for i in range(n):
|
||||
avg_dist = sum(distances[i]) / (n - 1) # Don't include distance to self
|
||||
avg_distances.append(avg_dist)
|
||||
|
||||
# Return the index of the text with the minimum average distance
|
||||
min_avg_dist = min(avg_distances)
|
||||
return avg_distances.index(min_avg_dist)
|
||||
|
||||
|
||||
def group_repeats(md_files: List[str]) -> Dict[str, List[str]]:
|
||||
"""
|
||||
Group MD files by their base name (without the repeat number).
|
||||
Returns a dictionary mapping base names to lists of file paths.
|
||||
"""
|
||||
grouped = {}
|
||||
|
||||
for md_path in md_files:
|
||||
base_name = re.sub(r"_repeat\d+\.md$", "", os.path.basename(md_path))
|
||||
if base_name not in grouped:
|
||||
grouped[base_name] = []
|
||||
grouped[base_name].append(md_path)
|
||||
|
||||
return grouped
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Find mediod (most representative) examples from repeated OCR outputs.")
|
||||
parser.add_argument(
|
||||
"--input", type=str, required=True, help="Path to the directory containing repeated OCR outputs (e.g., *_repeat1.md, *_repeat2.md, etc.)"
|
||||
)
|
||||
parser.add_argument("--output", type=str, required=True, help="Path to the directory where mediod examples will be copied")
|
||||
parser.add_argument("--min_repeats", type=int, default=3, help="Minimum number of repeats required to compute a mediod (default: 3)")
|
||||
args = parser.parse_args()
|
||||
|
||||
input_dir = args.input
|
||||
output_dir = args.output
|
||||
min_repeats = args.min_repeats
|
||||
|
||||
# Create output directory if it doesn't exist
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Find all markdown files in the input directory (recursive)
|
||||
md_files = glob.glob(os.path.join(input_dir, "**/*.md"), recursive=True)
|
||||
|
||||
if not md_files:
|
||||
print(f"No markdown files found in {input_dir}")
|
||||
return
|
||||
|
||||
# Group files by their base name
|
||||
grouped_files = group_repeats(md_files)
|
||||
|
||||
# Process each group
|
||||
successful = 0
|
||||
skipped = 0
|
||||
|
||||
print(f"Found {len(grouped_files)} unique test cases with repeats")
|
||||
|
||||
for base_name, file_paths in tqdm(grouped_files.items(), desc="Processing test cases"):
|
||||
# Skip if there aren't enough repeats
|
||||
if len(file_paths) < min_repeats:
|
||||
print(f"Skipping {base_name}: only {len(file_paths)} repeats (minimum {min_repeats} required)")
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
# Read all text content
|
||||
texts = []
|
||||
for path in file_paths:
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
texts.append(f.read())
|
||||
except Exception as e:
|
||||
print(f"Error reading {path}: {e}")
|
||||
continue
|
||||
|
||||
# Find the mediod
|
||||
mediod_idx = find_mediod(texts)
|
||||
if mediod_idx == -1:
|
||||
print(f"Failed to find mediod for {base_name}")
|
||||
skipped += 1
|
||||
continue
|
||||
|
||||
# Get the path of the mediod file
|
||||
mediod_path = file_paths[mediod_idx]
|
||||
|
||||
# Create the output path, preserving the directory structure relative to input_dir
|
||||
rel_path = os.path.relpath(mediod_path, input_dir)
|
||||
# Change the repeat number to 1 in the output filename
|
||||
output_filename = re.sub(r"_repeat\d+\.md$", "_repeat1.md", os.path.basename(rel_path))
|
||||
output_subdir = os.path.dirname(rel_path)
|
||||
output_path = os.path.join(output_dir, output_subdir, output_filename)
|
||||
|
||||
# Create directories if needed
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
# Copy the mediod file
|
||||
try:
|
||||
shutil.copy2(mediod_path, output_path)
|
||||
successful += 1
|
||||
except Exception as e:
|
||||
print(f"Error copying {mediod_path} to {output_path}: {e}")
|
||||
|
||||
print(f"Processing complete: {successful} mediods copied, {skipped} cases skipped")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,62 @@
|
||||
def build_basic_prompt() -> str:
|
||||
return "Please provide a natural, plain text representation of the document, formatted in Markdown. Skip any headers and footers. For ALL mathematical expressions, use LaTeX notation with \( and \) for inline equations and \[ and \] for display equations. Convert any tables into Markdown format."
|
||||
|
||||
|
||||
def build_openai_silver_data_prompt_no_document_anchoring(_base_text: str) -> str:
|
||||
return (
|
||||
"Below is the image of one page of a PDF document. "
|
||||
"Just return the plain text representation of this document as if you were reading it naturally.\n"
|
||||
"Turn equations into a LaTeX representation, and tables into markdown format. Remove the headers and footers, but keep references and footnotes.\n"
|
||||
"Read any natural handwriting.\n"
|
||||
"This is likely one page out of several in the document, so be sure to preserve any sentences that come from the previous page, or continue onto the next page, exactly as they are.\n"
|
||||
"If there is no text at all that you think you should read, you can output null.\n"
|
||||
"Do not hallucinate."
|
||||
)
|
||||
|
||||
|
||||
def claude_response_format_schema() -> dict:
|
||||
return (
|
||||
{
|
||||
"name": "page_response",
|
||||
"description": "Extracts text from pdf's.",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"primary_language": {
|
||||
"type": ["string", "null"],
|
||||
"description": "The primary language of the text using two-letter codes or null if there is no text at all that you think you should read.",
|
||||
},
|
||||
"is_rotation_valid": {
|
||||
"type": "boolean",
|
||||
"description": "Is this page oriented correctly for reading? Answer only considering the textual content, do not factor in the rotation of any charts, tables, drawings, or figures.",
|
||||
},
|
||||
"rotation_correction": {
|
||||
"type": "integer",
|
||||
"description": "Indicates the degree of clockwise rotation needed if the page is not oriented correctly.",
|
||||
"enum": [0, 90, 180, 270],
|
||||
"default": 0,
|
||||
},
|
||||
"is_table": {
|
||||
"type": "boolean",
|
||||
"description": "Indicates if the majority of the page content is in tabular format.",
|
||||
},
|
||||
"is_diagram": {
|
||||
"type": "boolean",
|
||||
"description": "Indicates if the majority of the page content is a visual diagram.",
|
||||
},
|
||||
"natural_text": {
|
||||
"type": ["string", "null"],
|
||||
"description": "The natural text content extracted from the page.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"primary_language",
|
||||
"is_rotation_valid",
|
||||
"rotation_correction",
|
||||
"is_table",
|
||||
"is_diagram",
|
||||
"natural_text",
|
||||
],
|
||||
},
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,300 @@
|
||||
import glob
|
||||
import os
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64webp
|
||||
|
||||
from .tests import BasePDFTest
|
||||
|
||||
|
||||
def _filter_by_max_reports(
|
||||
test_results_by_candidate: Dict[str, Dict[str, Dict[int, List[Tuple[BasePDFTest, bool, str]]]]],
|
||||
test_to_jsonl: Dict[str, str],
|
||||
max_reports: int,
|
||||
) -> Dict[str, Dict[str, Dict[int, List[Tuple[BasePDFTest, bool, str]]]]]:
|
||||
"""Filter test results to include at most max_reports unique PDFs per .jsonl file."""
|
||||
filtered = {}
|
||||
for candidate, pdf_results in test_results_by_candidate.items():
|
||||
# Track which unique PDFs we've included per jsonl file
|
||||
jsonl_pdfs: Dict[str, set] = {}
|
||||
filtered_pdfs: Dict[str, Dict[int, List[Tuple[BasePDFTest, bool, str]]]] = {}
|
||||
|
||||
for pdf_name in sorted(pdf_results.keys()):
|
||||
pages = pdf_results[pdf_name]
|
||||
for page in sorted(pages.keys()):
|
||||
for test_tuple in pages[page]:
|
||||
test, passed, explanation = test_tuple
|
||||
jsonl_file = test_to_jsonl.get(test.id, "unknown")
|
||||
jsonl_pdfs.setdefault(jsonl_file, set())
|
||||
# If this PDF is new for this jsonl file, check the limit
|
||||
if pdf_name not in jsonl_pdfs[jsonl_file]:
|
||||
if len(jsonl_pdfs[jsonl_file]) >= max_reports:
|
||||
continue
|
||||
jsonl_pdfs[jsonl_file].add(pdf_name)
|
||||
filtered_pdfs.setdefault(pdf_name, {}).setdefault(page, []).append(test_tuple)
|
||||
|
||||
filtered[candidate] = filtered_pdfs
|
||||
return filtered
|
||||
|
||||
|
||||
def generate_html_report(
|
||||
test_results_by_candidate: Dict[str, Dict[str, Dict[int, List[Tuple[BasePDFTest, bool, str]]]]],
|
||||
pdf_folder: str,
|
||||
output_file: str,
|
||||
max_reports: Optional[int] = None,
|
||||
test_to_jsonl: Optional[Dict[str, str]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Generate a simple static HTML report of test results.
|
||||
|
||||
Args:
|
||||
test_results_by_candidate: Dictionary mapping candidate name to dictionary mapping PDF name to dictionary
|
||||
mapping page number to list of (test, passed, explanation) tuples.
|
||||
pdf_folder: Path to the folder containing PDF files.
|
||||
output_file: Path to the output HTML file.
|
||||
max_reports: If set, limit to at most N tests per .jsonl file in the report.
|
||||
test_to_jsonl: Dictionary mapping test IDs to their source jsonl filenames.
|
||||
"""
|
||||
# If max_reports is set, filter test_results_by_candidate to limit per jsonl file
|
||||
if max_reports is not None and test_to_jsonl is not None:
|
||||
test_results_by_candidate = _filter_by_max_reports(test_results_by_candidate, test_to_jsonl, max_reports)
|
||||
candidates = list(test_results_by_candidate.keys())
|
||||
|
||||
# Create HTML report
|
||||
html = """<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>OLMOCR Bench Test Report</title>
|
||||
<style>
|
||||
body {
|
||||
font-family: Arial, sans-serif;
|
||||
margin: 0;
|
||||
padding: 20px;
|
||||
color: #333;
|
||||
line-height: 1.6;
|
||||
}
|
||||
|
||||
h1, h2, h3, h4 {
|
||||
margin-top: 20px;
|
||||
margin-bottom: 10px;
|
||||
}
|
||||
|
||||
.test-block {
|
||||
border: 1px solid #ddd;
|
||||
margin-bottom: 30px;
|
||||
padding: 15px;
|
||||
border-radius: 5px;
|
||||
}
|
||||
|
||||
.test-block.pass {
|
||||
border-left: 5px solid #4CAF50;
|
||||
}
|
||||
|
||||
.test-block.fail {
|
||||
border-left: 5px solid #F44336;
|
||||
}
|
||||
|
||||
.pdf-image {
|
||||
max-width: 100%;
|
||||
border: 1px solid #ddd;
|
||||
margin: 10px 0;
|
||||
}
|
||||
|
||||
.markdown-content {
|
||||
background: #f5f5f5;
|
||||
padding: 15px;
|
||||
border: 1px solid #ddd;
|
||||
border-radius: 5px;
|
||||
font-family: monospace;
|
||||
white-space: pre-wrap;
|
||||
overflow-x: auto;
|
||||
margin: 10px 0;
|
||||
}
|
||||
|
||||
.markdown-content table {
|
||||
border-collapse: collapse;
|
||||
margin: 10px 0;
|
||||
}
|
||||
|
||||
.markdown-content th,
|
||||
.markdown-content td {
|
||||
border: 1px solid #999;
|
||||
padding: 4px 8px;
|
||||
}
|
||||
|
||||
.markdown-content th {
|
||||
background-color: #e0e0e0;
|
||||
}
|
||||
|
||||
.status {
|
||||
display: inline-block;
|
||||
padding: 3px 8px;
|
||||
border-radius: 3px;
|
||||
font-weight: bold;
|
||||
margin-left: 10px;
|
||||
}
|
||||
|
||||
.pass-status {
|
||||
background-color: #4CAF50;
|
||||
color: white;
|
||||
}
|
||||
|
||||
.fail-status {
|
||||
background-color: #F44336;
|
||||
color: white;
|
||||
}
|
||||
|
||||
.test-details {
|
||||
margin: 10px 0;
|
||||
}
|
||||
|
||||
.test-explanation {
|
||||
margin-top: 10px;
|
||||
padding: 10px;
|
||||
background: #fff9c4;
|
||||
border-radius: 3px;
|
||||
}
|
||||
|
||||
hr {
|
||||
border: 0;
|
||||
border-top: 1px solid #ddd;
|
||||
margin: 30px 0;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>OLMOCR Bench Test Report</h1>
|
||||
"""
|
||||
|
||||
# Process all candidates
|
||||
print("Generating test report...")
|
||||
for candidate in candidates:
|
||||
html += f"<h2>Candidate: {candidate}</h2>\n"
|
||||
|
||||
# Get all PDFs for this candidate
|
||||
all_pdfs = sorted(test_results_by_candidate[candidate].keys())
|
||||
|
||||
for pdf_name in tqdm(all_pdfs, desc=f"Processing {candidate}"):
|
||||
pages = sorted(test_results_by_candidate[candidate][pdf_name].keys())
|
||||
|
||||
for page in pages:
|
||||
# Get tests for this PDF page
|
||||
tests = test_results_by_candidate[candidate][pdf_name][page]
|
||||
|
||||
for test, passed, explanation in tests:
|
||||
result_class = "pass" if passed else "fail"
|
||||
status_text = "PASSED" if passed else "FAILED"
|
||||
status_class = "pass-status" if passed else "fail-status"
|
||||
|
||||
# Begin test block
|
||||
html += f"""
|
||||
<div class="test-block {result_class}">
|
||||
<h3>Test ID: {test.id} <span class="status {status_class}">{status_text}</span></h3>
|
||||
<p><strong>PDF:</strong> {pdf_name} | <strong>Page:</strong> {page} | <strong>Type:</strong> {test.type}</p>
|
||||
|
||||
<div class="test-details">
|
||||
"""
|
||||
|
||||
# Add test details based on type
|
||||
test_type = getattr(test, "type", "").lower()
|
||||
if test_type == "present" and hasattr(test, "text"):
|
||||
text = getattr(test, "text", "")
|
||||
html += f""" <p><strong>Text to find:</strong> "{text}"</p>\n"""
|
||||
elif test_type == "absent" and hasattr(test, "text"):
|
||||
text = getattr(test, "text", "")
|
||||
html += f""" <p><strong>Text should not appear:</strong> "{text}"</p>\n"""
|
||||
elif test_type == "order" and hasattr(test, "before") and hasattr(test, "after"):
|
||||
before = getattr(test, "before", "")
|
||||
after = getattr(test, "after", "")
|
||||
html += f""" <p><strong>Text order:</strong> "{before}" should appear before "{after}"</p>\n"""
|
||||
elif test_type == "table":
|
||||
if hasattr(test, "cell"):
|
||||
cell = getattr(test, "cell", "")
|
||||
html += f""" <p><strong>Table cell:</strong> "{cell}"</p>\n"""
|
||||
if hasattr(test, "up") and getattr(test, "up", None):
|
||||
up = getattr(test, "up")
|
||||
html += f""" <p><strong>Above:</strong> "{up}"</p>\n"""
|
||||
if hasattr(test, "down") and getattr(test, "down", None):
|
||||
down = getattr(test, "down")
|
||||
html += f""" <p><strong>Below:</strong> "{down}"</p>\n"""
|
||||
if hasattr(test, "left") and getattr(test, "left", None):
|
||||
left = getattr(test, "left")
|
||||
html += f""" <p><strong>Left:</strong> "{left}"</p>\n"""
|
||||
if hasattr(test, "right") and getattr(test, "right", None):
|
||||
right = getattr(test, "right")
|
||||
html += f""" <p><strong>Right:</strong> "{right}"</p>\n"""
|
||||
elif test_type == "math" and hasattr(test, "math"):
|
||||
math = getattr(test, "math", "")
|
||||
html += f""" <p><strong>Math equation:</strong> {math}</p>\n"""
|
||||
elif test_type == "format" and hasattr(test, "text"):
|
||||
text = getattr(test, "text", "")
|
||||
fmt = getattr(test, "format", "")
|
||||
html += f""" <p><strong>Text:</strong> "{text}" should be formatted as <strong>{fmt}</strong></p>\n"""
|
||||
elif test_type == "footnote" and hasattr(test, "marker"):
|
||||
marker = getattr(test, "marker", "")
|
||||
html += f""" <p><strong>Footnote marker:</strong> {marker}</p>\n"""
|
||||
before = getattr(test, "appears_before_marker", None)
|
||||
after = getattr(test, "appears_after_marker", None)
|
||||
if before:
|
||||
html += f""" <p><strong>Text before marker:</strong> "{before}"</p>\n"""
|
||||
if after:
|
||||
html += f""" <p><strong>Text after marker:</strong> "{after}"</p>\n"""
|
||||
elif test_type == "baseline":
|
||||
max_length = getattr(test, "max_length", None)
|
||||
if max_length is not None:
|
||||
html += f""" <p><strong>Baseline check:</strong> max length {max_length} (blank page check)</p>\n"""
|
||||
else:
|
||||
html += f""" <p><strong>Baseline check:</strong> non-blank, no repeats, valid characters</p>\n"""
|
||||
|
||||
html += """ </div>\n"""
|
||||
|
||||
# Add explanation for failed tests
|
||||
if not passed:
|
||||
html += f""" <div class="test-explanation">
|
||||
<strong>Explanation:</strong> {explanation}
|
||||
</div>\n"""
|
||||
|
||||
# Render PDF page
|
||||
pdf_path = os.path.join(pdf_folder, pdf_name)
|
||||
try:
|
||||
html += """ <h4>PDF Render:</h4>\n"""
|
||||
image_data = render_pdf_to_base64webp(pdf_path, page, 1024)
|
||||
html += f""" <img class="pdf-image" alt="PDF Page {page}" src="data:image/webp;base64,{image_data}" />\n"""
|
||||
except Exception as e:
|
||||
html += f""" <p>Error rendering PDF: {str(e)}</p>\n"""
|
||||
|
||||
# Get the Markdown content for this page
|
||||
md_content = None
|
||||
try:
|
||||
md_base = os.path.splitext(pdf_name)[0]
|
||||
md_files = list(glob.glob(os.path.join(os.path.dirname(pdf_folder), candidate, f"{md_base}_pg{page}_repeat*.md")))
|
||||
if md_files:
|
||||
md_file_path = md_files[0] # Use the first repeat as an example
|
||||
with open(md_file_path, "r", encoding="utf-8") as f:
|
||||
md_content = f.read()
|
||||
except Exception as e:
|
||||
md_content = f"Error loading Markdown content: {str(e)}"
|
||||
|
||||
if md_content:
|
||||
html += """ <h4>Markdown Content:</h4>\n"""
|
||||
html += f""" <div class="markdown-content">{md_content}</div>\n"""
|
||||
|
||||
# End test block
|
||||
html += """ </div>\n"""
|
||||
|
||||
# Add separator between pages
|
||||
html += """ <hr>\n"""
|
||||
|
||||
# Close HTML
|
||||
html += """</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
with open(output_file, "w", encoding="utf-8") as f:
|
||||
f.write(html)
|
||||
|
||||
print(f"Simple HTML report generated: {output_file}")
|
||||
@@ -0,0 +1,300 @@
|
||||
#!/usr/bin/env python3
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from flask import Flask, jsonify, redirect, render_template, request, send_file, url_for
|
||||
|
||||
app = Flask(__name__)
|
||||
# Add static folder for KaTeX files
|
||||
app.static_folder = os.path.join(os.path.dirname(os.path.abspath(__file__)), "katex")
|
||||
|
||||
# Global state
|
||||
DATASET_DIR = ""
|
||||
DATASET_FILE = None
|
||||
CURRENT_PDF = None
|
||||
PDF_TESTS = {}
|
||||
ALL_PDFS = []
|
||||
FORCE = False # New global flag
|
||||
|
||||
|
||||
def find_next_unchecked_pdf() -> Optional[str]:
|
||||
"""Find the next PDF with at least one unchecked test."""
|
||||
global PDF_TESTS, ALL_PDFS
|
||||
|
||||
for pdf_name in ALL_PDFS:
|
||||
pdf_tests = PDF_TESTS[pdf_name]
|
||||
for test in pdf_tests:
|
||||
if test.get("checked") is None:
|
||||
return pdf_name
|
||||
return None
|
||||
|
||||
|
||||
def calculate_stats() -> dict:
|
||||
"""Calculate statistics for all tests in the dataset."""
|
||||
global PDF_TESTS
|
||||
|
||||
total_tests = 0
|
||||
null_status = 0
|
||||
verified_status = 0
|
||||
rejected_status = 0
|
||||
|
||||
for pdf_tests in PDF_TESTS.values():
|
||||
total_tests += len(pdf_tests)
|
||||
|
||||
for test in pdf_tests:
|
||||
status = test.get("checked")
|
||||
if status is None:
|
||||
null_status += 1
|
||||
elif status == "verified":
|
||||
verified_status += 1
|
||||
elif status == "rejected":
|
||||
rejected_status += 1
|
||||
|
||||
completion = 0
|
||||
if total_tests > 0:
|
||||
completion = (verified_status + rejected_status) / total_tests * 100
|
||||
|
||||
return {"total": total_tests, "null": null_status, "verified": verified_status, "rejected": rejected_status, "completion": completion}
|
||||
|
||||
|
||||
def save_dataset(jsonl_file: str) -> None:
|
||||
"""Save the tests to a JSONL file, using temp file for atomic write."""
|
||||
global PDF_TESTS
|
||||
|
||||
# Flatten all tests
|
||||
all_tests = []
|
||||
for pdf_tests in PDF_TESTS.values():
|
||||
all_tests.extend(pdf_tests)
|
||||
|
||||
# Create temp file and write updated content
|
||||
with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
|
||||
for test in all_tests:
|
||||
temp_file.write(json.dumps(test) + "\n")
|
||||
|
||||
# Atomic replace
|
||||
shutil.move(temp_file.name, jsonl_file)
|
||||
|
||||
|
||||
@app.route("/pdf/<path:pdf_name>")
|
||||
def serve_pdf(pdf_name):
|
||||
"""Serve the PDF file directly."""
|
||||
pdf_path = os.path.join(DATASET_DIR, "pdfs", pdf_name)
|
||||
return send_file(pdf_path, mimetype="application/pdf")
|
||||
|
||||
|
||||
@app.route("/")
|
||||
def index():
|
||||
"""Main page displaying the current PDF and its tests."""
|
||||
global CURRENT_PDF, PDF_TESTS, DATASET_DIR, ALL_PDFS, FORCE
|
||||
|
||||
# If no current PDF is set, find the next one with unchecked tests
|
||||
if CURRENT_PDF is None:
|
||||
CURRENT_PDF = find_next_unchecked_pdf()
|
||||
|
||||
# If still no PDF, either show the "All done" page or force display the first PDF
|
||||
if CURRENT_PDF is None:
|
||||
if FORCE and ALL_PDFS:
|
||||
CURRENT_PDF = ALL_PDFS[0]
|
||||
else:
|
||||
return render_template("all_done.html")
|
||||
|
||||
# Get the tests for the current PDF
|
||||
current_tests = PDF_TESTS.get(CURRENT_PDF, [])
|
||||
|
||||
# Create PDF URL for pdf.js to load
|
||||
pdf_url = url_for("serve_pdf", pdf_name=CURRENT_PDF)
|
||||
|
||||
# Calculate statistics
|
||||
stats = calculate_stats()
|
||||
|
||||
return render_template(
|
||||
"review.html",
|
||||
pdf_name=CURRENT_PDF,
|
||||
tests=current_tests,
|
||||
pdf_path=pdf_url,
|
||||
pdf_index=ALL_PDFS.index(CURRENT_PDF) if CURRENT_PDF in ALL_PDFS else 0,
|
||||
total_pdfs=len(ALL_PDFS),
|
||||
stats=stats,
|
||||
)
|
||||
|
||||
|
||||
@app.route("/update_test", methods=["POST"])
|
||||
def update_test():
|
||||
"""API endpoint to update a test."""
|
||||
global PDF_TESTS, DATASET_DIR, DATASET_FILE
|
||||
|
||||
data = request.json
|
||||
pdf_name = data.get("pdf")
|
||||
test_id = data.get("id")
|
||||
field = data.get("field")
|
||||
value = data.get("value")
|
||||
|
||||
# Find and update the test
|
||||
for test in PDF_TESTS.get(pdf_name, []):
|
||||
if test.get("id") == test_id:
|
||||
test[field] = value
|
||||
break
|
||||
|
||||
# Save the updated tests
|
||||
save_dataset(DATASET_FILE)
|
||||
|
||||
return jsonify({"status": "success"})
|
||||
|
||||
|
||||
@app.route("/reject_all", methods=["POST"])
|
||||
def reject_all():
|
||||
"""API endpoint to reject all tests for a PDF."""
|
||||
global PDF_TESTS, DATASET_DIR, DATASET_FILE
|
||||
|
||||
data = request.json
|
||||
pdf_name = data.get("pdf")
|
||||
|
||||
if pdf_name and pdf_name in PDF_TESTS:
|
||||
# Update all tests for this PDF to rejected
|
||||
for test in PDF_TESTS[pdf_name]:
|
||||
test["checked"] = "rejected"
|
||||
|
||||
# Save the updated tests
|
||||
save_dataset(DATASET_FILE)
|
||||
|
||||
return jsonify({"status": "success", "count": len(PDF_TESTS[pdf_name])})
|
||||
|
||||
return jsonify({"status": "error", "message": "PDF not found"})
|
||||
|
||||
|
||||
@app.route("/next_pdf", methods=["POST"])
|
||||
def next_pdf():
|
||||
"""Move to the next PDF in the list."""
|
||||
global CURRENT_PDF, ALL_PDFS, FORCE
|
||||
|
||||
if CURRENT_PDF in ALL_PDFS:
|
||||
current_index = ALL_PDFS.index(CURRENT_PDF)
|
||||
if current_index < len(ALL_PDFS) - 1:
|
||||
CURRENT_PDF = ALL_PDFS[current_index + 1]
|
||||
else:
|
||||
# If in force mode, cycle back to the beginning instead of checking for an unchecked PDF
|
||||
if FORCE and ALL_PDFS:
|
||||
CURRENT_PDF = ALL_PDFS[0]
|
||||
else:
|
||||
CURRENT_PDF = find_next_unchecked_pdf()
|
||||
else:
|
||||
if FORCE and ALL_PDFS:
|
||||
CURRENT_PDF = ALL_PDFS[0]
|
||||
else:
|
||||
CURRENT_PDF = find_next_unchecked_pdf()
|
||||
|
||||
return redirect(url_for("index"))
|
||||
|
||||
|
||||
@app.route("/prev_pdf", methods=["POST"])
|
||||
def prev_pdf():
|
||||
"""Move to the previous PDF in the list."""
|
||||
global CURRENT_PDF, ALL_PDFS
|
||||
|
||||
if CURRENT_PDF in ALL_PDFS:
|
||||
current_index = ALL_PDFS.index(CURRENT_PDF)
|
||||
if current_index > 0:
|
||||
CURRENT_PDF = ALL_PDFS[current_index - 1]
|
||||
|
||||
return redirect(url_for("index"))
|
||||
|
||||
|
||||
@app.route("/goto_pdf/<int:index>", methods=["POST"])
|
||||
def goto_pdf(index):
|
||||
"""Go to a specific PDF by index."""
|
||||
global CURRENT_PDF, ALL_PDFS
|
||||
|
||||
if 0 <= index < len(ALL_PDFS):
|
||||
CURRENT_PDF = ALL_PDFS[index]
|
||||
|
||||
return redirect(url_for("index"))
|
||||
|
||||
|
||||
def load_dataset(dataset_file: str) -> Tuple[Dict[str, List[Dict]], List[str]]:
|
||||
"""Load tests from the dataset file and organize them by PDF."""
|
||||
if not os.path.exists(dataset_file):
|
||||
raise FileNotFoundError(f"Dataset file not found: {dataset_file}")
|
||||
|
||||
pdf_tests = defaultdict(list)
|
||||
|
||||
with open(dataset_file, "r") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
try:
|
||||
test = json.loads(line)
|
||||
pdf_name = test.get("pdf")
|
||||
if pdf_name:
|
||||
pdf_tests[pdf_name].append(test)
|
||||
except json.JSONDecodeError:
|
||||
print(f"Warning: Could not parse line as JSON: {line}")
|
||||
|
||||
all_pdfs = list(pdf_tests.keys())
|
||||
|
||||
return pdf_tests, all_pdfs
|
||||
|
||||
|
||||
def create_templates_directory():
|
||||
"""Create templates directory for Flask if it doesn't exist."""
|
||||
templates_dir = os.path.join(os.path.dirname(__file__), "templates")
|
||||
os.makedirs(templates_dir, exist_ok=True)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point with command-line arguments."""
|
||||
global DATASET_DIR, DATASET_FILE, PDF_TESTS, ALL_PDFS, CURRENT_PDF, FORCE
|
||||
|
||||
parser = argparse.ArgumentParser(description="Interactive Test Review App")
|
||||
parser.add_argument("dataset_file", help="Path to the dataset jsonl file")
|
||||
parser.add_argument("--port", type=int, default=5000, help="Port for the Flask app")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Host for the Flask app")
|
||||
parser.add_argument("--debug", action="store_true", help="Run Flask in debug mode")
|
||||
parser.add_argument("--force", action="store_true", help="Force show each file one by one and never do the 'All done' page")
|
||||
|
||||
args = parser.parse_args()
|
||||
FORCE = args.force # Set the global FORCE flag
|
||||
|
||||
# Validate dataset directory
|
||||
if not os.path.exists(args.dataset_file):
|
||||
print(f"Error: Dataset not found: {args.dataset_file}")
|
||||
return 1
|
||||
|
||||
# Store dataset directory globally
|
||||
DATASET_DIR = os.path.dirname(os.path.abspath(args.dataset_file))
|
||||
DATASET_FILE = args.dataset_file
|
||||
|
||||
pdf_dir = os.path.join(DATASET_DIR, "pdfs")
|
||||
if not os.path.isdir(pdf_dir):
|
||||
print(f"Error: PDF directory not found: {pdf_dir}")
|
||||
return 1
|
||||
|
||||
# Load dataset
|
||||
try:
|
||||
PDF_TESTS, ALL_PDFS = load_dataset(args.dataset_file)
|
||||
except Exception as e:
|
||||
print(f"Error loading dataset: {str(e)}")
|
||||
return 1
|
||||
|
||||
# Create templates directory
|
||||
create_templates_directory()
|
||||
|
||||
# Find first PDF with unchecked tests
|
||||
CURRENT_PDF = find_next_unchecked_pdf()
|
||||
|
||||
# Start Flask app
|
||||
print(f"Starting server at http://{args.host}:{args.port}")
|
||||
app.run(host=args.host, port=args.port, debug=args.debug)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,716 @@
|
||||
#!/usr/bin/env python3
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from flask import Flask, jsonify, redirect, render_template, request, send_file, url_for
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
# Global state
|
||||
DATASET_DIR = ""
|
||||
DATASET_FILE = None
|
||||
CURRENT_PDF = None
|
||||
PDF_TESTS = {}
|
||||
ALL_PDFS = []
|
||||
FORCE = False
|
||||
|
||||
|
||||
def find_next_unchecked_pdf() -> Optional[str]:
|
||||
"""Find the next PDF with at least one unchecked test."""
|
||||
global PDF_TESTS, ALL_PDFS
|
||||
|
||||
for pdf_name in ALL_PDFS:
|
||||
pdf_tests = PDF_TESTS[pdf_name]
|
||||
for test in pdf_tests:
|
||||
if test.get("checked") is None:
|
||||
return pdf_name
|
||||
return None
|
||||
|
||||
|
||||
def calculate_stats() -> dict:
|
||||
"""Calculate statistics for all tests in the dataset."""
|
||||
global PDF_TESTS
|
||||
|
||||
total_tests = 0
|
||||
null_status = 0
|
||||
verified_status = 0
|
||||
rejected_status = 0
|
||||
|
||||
for pdf_tests in PDF_TESTS.values():
|
||||
total_tests += len(pdf_tests)
|
||||
|
||||
for test in pdf_tests:
|
||||
status = test.get("checked")
|
||||
if status is None:
|
||||
null_status += 1
|
||||
elif status == "verified":
|
||||
verified_status += 1
|
||||
elif status == "rejected":
|
||||
rejected_status += 1
|
||||
|
||||
completion = 0
|
||||
if total_tests > 0:
|
||||
completion = (verified_status + rejected_status) / total_tests * 100
|
||||
|
||||
return {"total": total_tests, "null": null_status, "verified": verified_status, "rejected": rejected_status, "completion": completion}
|
||||
|
||||
|
||||
def save_dataset(jsonl_file: str) -> None:
|
||||
"""Save the tests to a JSONL file, using temp file for atomic write."""
|
||||
global PDF_TESTS
|
||||
|
||||
# Flatten all tests
|
||||
all_tests = []
|
||||
for pdf_tests in PDF_TESTS.values():
|
||||
all_tests.extend(pdf_tests)
|
||||
|
||||
# Create temp file and write updated content
|
||||
with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
|
||||
for test in all_tests:
|
||||
temp_file.write(json.dumps(test) + "\n")
|
||||
|
||||
# Atomic replace
|
||||
shutil.move(temp_file.name, jsonl_file)
|
||||
|
||||
|
||||
@app.route("/pdf/<path:pdf_name>")
|
||||
def serve_pdf(pdf_name):
|
||||
"""Serve the PDF file directly."""
|
||||
pdf_path = os.path.join(DATASET_DIR, "pdfs", pdf_name)
|
||||
return send_file(pdf_path, mimetype="application/pdf")
|
||||
|
||||
|
||||
@app.route("/")
|
||||
def index():
|
||||
"""Main page displaying the current PDF and its tests."""
|
||||
global CURRENT_PDF, PDF_TESTS, DATASET_DIR, ALL_PDFS, FORCE
|
||||
|
||||
# If no current PDF is set, find the next one with unchecked tests
|
||||
if CURRENT_PDF is None:
|
||||
CURRENT_PDF = find_next_unchecked_pdf()
|
||||
|
||||
# If still no PDF, either show the "All done" page or force display the first PDF
|
||||
if CURRENT_PDF is None:
|
||||
if FORCE and ALL_PDFS:
|
||||
CURRENT_PDF = ALL_PDFS[0]
|
||||
else:
|
||||
return render_template("all_done_latex.html")
|
||||
|
||||
# Get the tests for the current PDF
|
||||
current_tests = PDF_TESTS.get(CURRENT_PDF, [])
|
||||
|
||||
# Create PDF URL for pdf.js to load
|
||||
pdf_url = url_for("serve_pdf", pdf_name=CURRENT_PDF)
|
||||
|
||||
# Calculate statistics
|
||||
stats = calculate_stats()
|
||||
|
||||
return render_template(
|
||||
"review_latex.html",
|
||||
pdf_name=CURRENT_PDF,
|
||||
tests=current_tests,
|
||||
pdf_path=pdf_url,
|
||||
pdf_index=ALL_PDFS.index(CURRENT_PDF) if CURRENT_PDF in ALL_PDFS else 0,
|
||||
total_pdfs=len(ALL_PDFS),
|
||||
stats=stats,
|
||||
)
|
||||
|
||||
|
||||
@app.route("/update_test", methods=["POST"])
|
||||
def update_test():
|
||||
"""API endpoint to update a test."""
|
||||
global PDF_TESTS, DATASET_DIR, DATASET_FILE
|
||||
|
||||
data = request.json
|
||||
pdf_name = data.get("pdf")
|
||||
test_id = data.get("id")
|
||||
field = data.get("field")
|
||||
value = data.get("value")
|
||||
|
||||
# Find and update the test
|
||||
for test in PDF_TESTS.get(pdf_name, []):
|
||||
if test.get("id") == test_id:
|
||||
test[field] = value
|
||||
break
|
||||
|
||||
# Save the updated tests
|
||||
save_dataset(DATASET_FILE)
|
||||
|
||||
return jsonify({"status": "success"})
|
||||
|
||||
|
||||
@app.route("/reject_all", methods=["POST"])
|
||||
def reject_all():
|
||||
"""API endpoint to reject all tests for a PDF."""
|
||||
global PDF_TESTS, DATASET_DIR, DATASET_FILE
|
||||
|
||||
data = request.json
|
||||
pdf_name = data.get("pdf")
|
||||
|
||||
if pdf_name and pdf_name in PDF_TESTS:
|
||||
# Update all tests for this PDF to rejected
|
||||
for test in PDF_TESTS[pdf_name]:
|
||||
test["checked"] = "rejected"
|
||||
|
||||
# Save the updated tests
|
||||
save_dataset(DATASET_FILE)
|
||||
|
||||
return jsonify({"status": "success", "count": len(PDF_TESTS[pdf_name])})
|
||||
|
||||
return jsonify({"status": "error", "message": "PDF not found"})
|
||||
|
||||
|
||||
@app.route("/next_pdf", methods=["POST"])
|
||||
def next_pdf():
|
||||
"""Move to the next PDF in the list."""
|
||||
global CURRENT_PDF, ALL_PDFS, FORCE
|
||||
|
||||
if CURRENT_PDF in ALL_PDFS:
|
||||
current_index = ALL_PDFS.index(CURRENT_PDF)
|
||||
if current_index < len(ALL_PDFS) - 1:
|
||||
CURRENT_PDF = ALL_PDFS[current_index + 1]
|
||||
else:
|
||||
# If in force mode, cycle back to the beginning instead of checking for an unchecked PDF
|
||||
if FORCE and ALL_PDFS:
|
||||
CURRENT_PDF = ALL_PDFS[0]
|
||||
else:
|
||||
CURRENT_PDF = find_next_unchecked_pdf()
|
||||
else:
|
||||
if FORCE and ALL_PDFS:
|
||||
CURRENT_PDF = ALL_PDFS[0]
|
||||
else:
|
||||
CURRENT_PDF = find_next_unchecked_pdf()
|
||||
|
||||
return redirect(url_for("index"))
|
||||
|
||||
|
||||
@app.route("/prev_pdf", methods=["POST"])
|
||||
def prev_pdf():
|
||||
"""Move to the previous PDF in the list."""
|
||||
global CURRENT_PDF, ALL_PDFS
|
||||
|
||||
if CURRENT_PDF in ALL_PDFS:
|
||||
current_index = ALL_PDFS.index(CURRENT_PDF)
|
||||
if current_index > 0:
|
||||
CURRENT_PDF = ALL_PDFS[current_index - 1]
|
||||
|
||||
return redirect(url_for("index"))
|
||||
|
||||
|
||||
@app.route("/goto_pdf/<int:index>", methods=["POST"])
|
||||
def goto_pdf(index):
|
||||
"""Go to a specific PDF by index."""
|
||||
global CURRENT_PDF, ALL_PDFS
|
||||
|
||||
if 0 <= index < len(ALL_PDFS):
|
||||
CURRENT_PDF = ALL_PDFS[index]
|
||||
|
||||
return redirect(url_for("index"))
|
||||
|
||||
|
||||
def load_dataset(dataset_file: str) -> Tuple[Dict[str, List[Dict]], List[str]]:
|
||||
"""Load tests from the dataset file and organize them by PDF."""
|
||||
if not os.path.exists(dataset_file):
|
||||
raise FileNotFoundError(f"Dataset file not found: {dataset_file}")
|
||||
|
||||
pdf_tests = defaultdict(list)
|
||||
|
||||
with open(dataset_file, "r") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
try:
|
||||
test = json.loads(line)
|
||||
pdf_name = test.get("pdf")
|
||||
if pdf_name:
|
||||
pdf_tests[pdf_name].append(test)
|
||||
except json.JSONDecodeError:
|
||||
print(f"Warning: Could not parse line as JSON: {line}")
|
||||
|
||||
all_pdfs = list(pdf_tests.keys())
|
||||
|
||||
return pdf_tests, all_pdfs
|
||||
|
||||
|
||||
def create_templates_directory():
|
||||
"""Create templates directory for Flask if it doesn't exist."""
|
||||
templates_dir = os.path.join(os.path.dirname(__file__), "templates")
|
||||
os.makedirs(templates_dir, exist_ok=True)
|
||||
|
||||
# Create the review_latex.html template with MathJax support
|
||||
review_html = """
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<!-- You can adjust the viewport settings as needed -->
|
||||
<meta name="viewport" content="width=1200, initial-scale=1.0">
|
||||
<title>Equation Verification</title>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/pdf.js/3.4.120/pdf.min.js"></script>
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/3.2.2/es5/tex-mml-chtml.js"></script>
|
||||
<style>
|
||||
body {
|
||||
font-family: Arial, sans-serif;
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
height: 100vh;
|
||||
}
|
||||
.header {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
align-items: center;
|
||||
padding: 10px;
|
||||
background-color: #f0f0f0;
|
||||
border-bottom: 1px solid #ddd;
|
||||
}
|
||||
.content {
|
||||
display: flex;
|
||||
flex: 1;
|
||||
overflow: hidden;
|
||||
}
|
||||
.pdf-viewer {
|
||||
flex: 2; /* Increased from 1 to 2 to make PDF larger */
|
||||
border-right: 1px solid #ddd;
|
||||
overflow: hidden;
|
||||
position: relative;
|
||||
}
|
||||
/* Updated PDF container size */
|
||||
#pdf-container {
|
||||
width: 200%; /* New fixed width */
|
||||
height: 200%; /* New fixed height */
|
||||
overflow: auto;
|
||||
}
|
||||
#zoom-controls {
|
||||
position: absolute;
|
||||
top: 10px;
|
||||
right: 10px;
|
||||
z-index: 100;
|
||||
background-color: white;
|
||||
padding: 5px;
|
||||
border-radius: 4px;
|
||||
box-shadow: 0 1px 3px rgba(0,0,0,0.2);
|
||||
}
|
||||
#zoom-controls button {
|
||||
margin: 0 5px;
|
||||
padding: 5px 10px;
|
||||
cursor: pointer;
|
||||
}
|
||||
.tests-panel {
|
||||
width: 1000px;
|
||||
overflow-y: auto;
|
||||
padding: 10px;
|
||||
}
|
||||
.test-item {
|
||||
border: 1px solid #ddd;
|
||||
border-radius: 4px;
|
||||
padding: 10px;
|
||||
margin-bottom: 10px;
|
||||
transition: background-color 0.2s;
|
||||
}
|
||||
.test-item.verified {
|
||||
background-color: #d4edda;
|
||||
}
|
||||
.test-item.rejected {
|
||||
background-color: #f8d7da;
|
||||
}
|
||||
/* The equation-display now stores the raw LaTeX in a data attribute */
|
||||
.equation-display {
|
||||
padding: 10px;
|
||||
margin: 5px 0;
|
||||
border: 1px solid #ddd;
|
||||
border-radius: 4px;
|
||||
background-color: #f9f9f9;
|
||||
overflow-x: auto;
|
||||
font-size: 1.2em; /* Larger font for equations */
|
||||
}
|
||||
.button-group {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
margin-top: 10px;
|
||||
}
|
||||
.button-group button {
|
||||
padding: 5px 10px;
|
||||
border: none;
|
||||
border-radius: 4px;
|
||||
cursor: pointer;
|
||||
}
|
||||
.verify-button {
|
||||
background-color: #28a745;
|
||||
color: white;
|
||||
}
|
||||
.reject-button {
|
||||
background-color: #dc3545;
|
||||
color: white;
|
||||
}
|
||||
.edit-button {
|
||||
background-color: #007bff;
|
||||
color: white;
|
||||
}
|
||||
.navigation {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
}
|
||||
.status {
|
||||
margin-left: 20px;
|
||||
}
|
||||
.progress-bar {
|
||||
height: 10px;
|
||||
background-color: #e9ecef;
|
||||
border-radius: 5px;
|
||||
margin-top: 5px;
|
||||
overflow: hidden;
|
||||
}
|
||||
.progress {
|
||||
height: 100%;
|
||||
background-color: #007bff;
|
||||
width: 0%;
|
||||
}
|
||||
/* Make MathJax equations more visible */
|
||||
.MathJax {
|
||||
font-size: 120% !important;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="header">
|
||||
<div>
|
||||
<h2>PDF: {{ pdf_name }}</h2>
|
||||
<form method="post" action="/reject_all" id="reject-all-form" style="display:inline;">
|
||||
<button type="button" onclick="rejectAll('{{ pdf_name }}')">Reject All Equations</button>
|
||||
</form>
|
||||
</div>
|
||||
<div class="navigation">
|
||||
<form method="post" action="/prev_pdf" style="display:inline;">
|
||||
<button type="submit" {% if pdf_index == 0 %}disabled{% endif %}>Previous</button>
|
||||
</form>
|
||||
<span style="margin: 0 10px;">{{ pdf_index + 1 }} / {{ total_pdfs }}</span>
|
||||
<form method="post" action="/next_pdf" style="display:inline;">
|
||||
<button type="submit">Next</button>
|
||||
</form>
|
||||
<div class="status">
|
||||
<div>Completion: {{ "%.1f"|format(stats.completion) }}%</div>
|
||||
<div class="progress-bar">
|
||||
<div class="progress" style="width: {{ stats.completion }}%;"></div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="content">
|
||||
<div class="pdf-viewer">
|
||||
<div id="zoom-controls">
|
||||
<button onclick="changeZoom(0.2)">+</button>
|
||||
<button onclick="changeZoom(-0.2)">-</button>
|
||||
<button onclick="resetZoom()">Reset</button>
|
||||
</div>
|
||||
<div id="pdf-container"></div>
|
||||
</div>
|
||||
<div class="tests-panel">
|
||||
<h3>Equations ({{ tests|length }})</h3>
|
||||
{% for test in tests %}
|
||||
<!-- Added data-latex attribute to store raw LaTeX -->
|
||||
<div class="test-item {% if test.checked == 'verified' %}verified{% elif test.checked == 'rejected' %}rejected{% endif %}" id="test-{{ test.id }}">
|
||||
<div class="equation-display" data-latex="{{ test.text|e }}">
|
||||
{{ test.text|safe }}
|
||||
</div>
|
||||
<div class="button-group">
|
||||
<button class="verify-button" onclick="updateTest('{{ test.id }}', '{{ test.pdf }}', 'checked', 'verified')">Verify</button>
|
||||
<button class="reject-button" onclick="updateTest('{{ test.id }}', '{{ test.pdf }}', 'checked', 'rejected')">Reject</button>
|
||||
<!-- New Edit button -->
|
||||
<button class="edit-button" onclick="enableEdit('{{ test.id }}', '{{ test.pdf }}')">Edit</button>
|
||||
</div>
|
||||
</div>
|
||||
{% endfor %}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
// Set up PDF.js
|
||||
pdfjsLib.GlobalWorkerOptions.workerSrc = 'https://cdnjs.cloudflare.com/ajax/libs/pdf.js/3.4.120/pdf.worker.min.js';
|
||||
|
||||
// Track current zoom level
|
||||
let currentScale = 2.0; // Initial larger scale
|
||||
let pdfDoc = null;
|
||||
let pageNum = 1;
|
||||
let canvas = null;
|
||||
|
||||
// Load the PDF
|
||||
const loadingTask = pdfjsLib.getDocument('{{ pdf_path }}');
|
||||
loadingTask.promise.then(function(pdf) {
|
||||
pdfDoc = pdf;
|
||||
renderPage(pageNum);
|
||||
});
|
||||
|
||||
// Function to render a page with the current scale
|
||||
function renderPage(num) {
|
||||
pdfDoc.getPage(num).then(function(page) {
|
||||
const viewport = page.getViewport({ scale: currentScale });
|
||||
if (!canvas) {
|
||||
canvas = document.createElement('canvas');
|
||||
document.getElementById('pdf-container').appendChild(canvas);
|
||||
}
|
||||
const context = canvas.getContext('2d');
|
||||
canvas.height = viewport.height;
|
||||
canvas.width = viewport.width;
|
||||
const renderContext = {
|
||||
canvasContext: context,
|
||||
viewport: viewport
|
||||
};
|
||||
page.render(renderContext);
|
||||
});
|
||||
}
|
||||
|
||||
// Function to change zoom level
|
||||
function changeZoom(delta) {
|
||||
currentScale += delta;
|
||||
if (currentScale < 0.5) currentScale = 0.5;
|
||||
if (currentScale > 5) currentScale = 5;
|
||||
renderPage(pageNum);
|
||||
}
|
||||
|
||||
// Function to reset zoom
|
||||
function resetZoom() {
|
||||
currentScale = 2.0;
|
||||
renderPage(pageNum);
|
||||
}
|
||||
|
||||
// Function to update a test – used by both verify/reject and edit
|
||||
function updateTest(testId, pdfName, field, value) {
|
||||
fetch('/update_test', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json'
|
||||
},
|
||||
body: JSON.stringify({
|
||||
id: testId,
|
||||
pdf: pdfName,
|
||||
field: field,
|
||||
value: value
|
||||
})
|
||||
})
|
||||
.then(response => response.json())
|
||||
.then(data => {
|
||||
if (data.status === 'success') {
|
||||
const testElement = document.getElementById(`test-${testId}`);
|
||||
testElement.classList.remove('verified', 'rejected');
|
||||
// Only update the class if the field updated is "checked".
|
||||
if (field === 'checked') {
|
||||
testElement.classList.add(value);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// New function to enable editing the LaTeX equation
|
||||
function enableEdit(testId, pdfName) {
|
||||
let testElement = document.getElementById("test-" + testId);
|
||||
let equationDisplay = testElement.querySelector(".equation-display");
|
||||
// Retrieve the raw LaTeX from the data attribute
|
||||
let rawLatex = equationDisplay.getAttribute('data-latex');
|
||||
// Save the current rendered HTML in case of cancellation
|
||||
let originalHTML = equationDisplay.innerHTML;
|
||||
|
||||
// Create a textarea for editing the LaTeX
|
||||
let textarea = document.createElement("textarea");
|
||||
textarea.id = "edit-input-" + testId;
|
||||
// Use the stored raw LaTeX if available; otherwise, fallback to textContent
|
||||
textarea.value = rawLatex ? rawLatex : equationDisplay.textContent.trim();
|
||||
textarea.style.width = "100%";
|
||||
textarea.rows = 3;
|
||||
|
||||
// Create Save button to commit changes
|
||||
let saveButton = document.createElement("button");
|
||||
saveButton.innerText = "Save";
|
||||
saveButton.onclick = function() {
|
||||
let newText = textarea.value;
|
||||
// Update the test via AJAX – updating the 'text' field
|
||||
updateTest(testId, pdfName, "text", newText);
|
||||
// Update the data-latex attribute to hold the new raw LaTeX code
|
||||
equationDisplay.setAttribute('data-latex', newText);
|
||||
// Replace the display content with the wrapped LaTeX for MathJax to process
|
||||
equationDisplay.innerHTML = '$$' + newText + '$$';
|
||||
if (typeof MathJax !== 'undefined') {
|
||||
MathJax.typeset();
|
||||
}
|
||||
// Clean up the temporary editing elements
|
||||
textarea.remove();
|
||||
saveButton.remove();
|
||||
cancelButton.remove();
|
||||
};
|
||||
|
||||
// Create Cancel button to revert changes
|
||||
let cancelButton = document.createElement("button");
|
||||
cancelButton.innerText = "Cancel";
|
||||
cancelButton.onclick = function() {
|
||||
equationDisplay.innerHTML = originalHTML;
|
||||
textarea.remove();
|
||||
saveButton.remove();
|
||||
cancelButton.remove();
|
||||
};
|
||||
|
||||
// Show the editing interface: clear the display and insert the textarea
|
||||
equationDisplay.innerHTML = "";
|
||||
equationDisplay.appendChild(textarea);
|
||||
// Append buttons to the test element (or you can choose to append them elsewhere)
|
||||
testElement.appendChild(saveButton);
|
||||
testElement.appendChild(cancelButton);
|
||||
}
|
||||
|
||||
// Function to reject all tests for a PDF
|
||||
function rejectAll(pdfName) {
|
||||
if (confirm('Are you sure you want to reject all equations for this PDF?')) {
|
||||
fetch('/reject_all', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json'
|
||||
},
|
||||
body: JSON.stringify({
|
||||
pdf: pdfName
|
||||
})
|
||||
})
|
||||
.then(response => response.json())
|
||||
.then(data => {
|
||||
if (data.status === 'success') {
|
||||
const testElements = document.querySelectorAll('.test-item');
|
||||
testElements.forEach(element => {
|
||||
element.classList.remove('verified');
|
||||
element.classList.add('rejected');
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Process LaTeX equations on page load: wrap plain text with $$ and trigger MathJax typesetting
|
||||
document.addEventListener('DOMContentLoaded', function() {
|
||||
const equationDisplays = document.querySelectorAll('.equation-display');
|
||||
equationDisplays.forEach(display => {
|
||||
let equation = display.textContent.trim();
|
||||
if (!equation.startsWith('$$')) {
|
||||
display.innerHTML = '$$' + equation + '$$';
|
||||
}
|
||||
});
|
||||
if (typeof MathJax !== 'undefined') {
|
||||
MathJax.typeset();
|
||||
}
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
# Create the all_done_latex.html template
|
||||
all_done_html = """
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>All Done!</title>
|
||||
<style>
|
||||
body {
|
||||
font-family: Arial, sans-serif;
|
||||
display: flex;
|
||||
justify-content: center;
|
||||
align-items: center;
|
||||
height: 100vh;
|
||||
margin: 0;
|
||||
background-color: #f5f5f5;
|
||||
}
|
||||
.container {
|
||||
text-align: center;
|
||||
padding: 30px;
|
||||
background-color: white;
|
||||
border-radius: 10px;
|
||||
box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
|
||||
}
|
||||
h1 {
|
||||
color: #28a745;
|
||||
margin-bottom: 20px;
|
||||
}
|
||||
p {
|
||||
font-size: 18px;
|
||||
margin-bottom: 20px;
|
||||
}
|
||||
button {
|
||||
padding: 10px 20px;
|
||||
background-color: #007bff;
|
||||
color: white;
|
||||
border: none;
|
||||
border-radius: 5px;
|
||||
cursor: pointer;
|
||||
font-size: 16px;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<h1>All Done! 🎉</h1>
|
||||
<p>You have reviewed all equations in the dataset.</p>
|
||||
<form method="post" action="/next_pdf">
|
||||
<button type="submit">Start Over</button>
|
||||
</form>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
with open(os.path.join(templates_dir, "review_latex.html"), "w") as f:
|
||||
f.write(review_html)
|
||||
|
||||
with open(os.path.join(templates_dir, "all_done_latex.html"), "w") as f:
|
||||
f.write(all_done_html)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point with command-line arguments."""
|
||||
global DATASET_DIR, DATASET_FILE, PDF_TESTS, ALL_PDFS, CURRENT_PDF, FORCE
|
||||
|
||||
parser = argparse.ArgumentParser(description="Interactive Test Review App")
|
||||
parser.add_argument("dataset_file", help="Path to the dataset jsonl file")
|
||||
parser.add_argument("--port", type=int, default=5000, help="Port for the Flask app")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Host for the Flask app")
|
||||
parser.add_argument("--debug", action="store_true", help="Run Flask in debug mode")
|
||||
parser.add_argument("--force", action="store_true", help="Force show each file one by one and never do the 'All done' page")
|
||||
|
||||
args = parser.parse_args()
|
||||
FORCE = args.force
|
||||
|
||||
if not os.path.exists(args.dataset_file):
|
||||
print(f"Error: Dataset not found: {args.dataset_file}")
|
||||
return 1
|
||||
|
||||
DATASET_DIR = os.path.dirname(os.path.abspath(args.dataset_file))
|
||||
DATASET_FILE = args.dataset_file
|
||||
|
||||
pdf_dir = os.path.join(DATASET_DIR, "pdfs")
|
||||
if not os.path.isdir(pdf_dir):
|
||||
print(f"Error: PDF directory not found: {pdf_dir}")
|
||||
return 1
|
||||
|
||||
try:
|
||||
PDF_TESTS, ALL_PDFS = load_dataset(args.dataset_file)
|
||||
except Exception as e:
|
||||
print(f"Error loading dataset: {str(e)}")
|
||||
return 1
|
||||
|
||||
create_templates_directory()
|
||||
CURRENT_PDF = find_next_unchecked_pdf()
|
||||
|
||||
print(f"Starting server at http://{args.host}:{args.port}")
|
||||
app.run(host=args.host, port=args.port, debug=args.debug)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,128 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from olmocr.bench.prompts import (
|
||||
build_basic_prompt,
|
||||
build_openai_silver_data_prompt_no_document_anchoring,
|
||||
)
|
||||
from olmocr.data.renderpdf import (
|
||||
get_png_dimensions_from_base64,
|
||||
render_pdf_to_base64png,
|
||||
)
|
||||
from olmocr.prompts.anchor import get_anchor_text
|
||||
from olmocr.prompts.prompts import (
|
||||
PageResponse,
|
||||
build_finetuning_prompt,
|
||||
build_openai_silver_data_prompt,
|
||||
build_openai_silver_data_prompt_v2,
|
||||
build_openai_silver_data_prompt_v2_simple,
|
||||
build_openai_silver_data_prompt_v3_simple,
|
||||
openai_response_format_schema,
|
||||
)
|
||||
|
||||
# Set up logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Global variables to track token usage and document count
|
||||
TOTAL_INPUT_TOKENS = 0
|
||||
TOTAL_OUTPUT_TOKENS = 0
|
||||
TOTAL_DOCUMENTS = 0
|
||||
|
||||
|
||||
def run_chatgpt(
|
||||
pdf_path: str,
|
||||
page_num: int = 1,
|
||||
model: str = "gpt-4o-2024-08-06",
|
||||
temperature: float = 0.1,
|
||||
target_longest_image_dim: int = 2048,
|
||||
max_completion_tokens: int = 10000,
|
||||
prompt_template: Literal["full", "full_no_document_anchoring", "basic", "finetune", "fullv2", "fullv2simple", "fullv3simple"] = "finetune",
|
||||
response_template: Literal["plain", "json"] = "json",
|
||||
) -> str:
|
||||
"""
|
||||
Convert page of a PDF file to markdown using the commercial openAI APIs.
|
||||
|
||||
See run_server.py for running against an openai compatible server
|
||||
|
||||
Args:
|
||||
pdf_path (str): The local path to the PDF file.
|
||||
|
||||
Returns:
|
||||
str: The OCR result in markdown format.
|
||||
"""
|
||||
global TOTAL_INPUT_TOKENS, TOTAL_OUTPUT_TOKENS, TOTAL_DOCUMENTS
|
||||
# Convert the first page of the PDF to a base64-encoded PNG image.
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=target_longest_image_dim)
|
||||
anchor_text = get_anchor_text(pdf_path, page_num, pdf_engine="pdfreport")
|
||||
|
||||
if not os.getenv("OPENAI_API_KEY"):
|
||||
raise SystemExit("You must specify an OPENAI_API_KEY")
|
||||
|
||||
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
|
||||
|
||||
if prompt_template == "full":
|
||||
prompt = build_openai_silver_data_prompt(anchor_text)
|
||||
elif prompt_template == "full_no_document_anchoring":
|
||||
prompt = build_openai_silver_data_prompt_no_document_anchoring(anchor_text)
|
||||
elif prompt_template == "finetune":
|
||||
prompt = build_finetuning_prompt(anchor_text)
|
||||
elif prompt_template == "basic":
|
||||
prompt = build_basic_prompt()
|
||||
elif prompt_template == "fullv2":
|
||||
prompt = build_openai_silver_data_prompt_v2(anchor_text)
|
||||
elif prompt_template == "fullv2simple":
|
||||
width, height = get_png_dimensions_from_base64(image_base64)
|
||||
prompt = build_openai_silver_data_prompt_v2_simple(width, height)
|
||||
elif prompt_template == "fullv3simple":
|
||||
width, height = get_png_dimensions_from_base64(image_base64)
|
||||
prompt = build_openai_silver_data_prompt_v3_simple(width, height)
|
||||
else:
|
||||
raise ValueError("Unknown prompt template")
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
temperature=temperature,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
# reasoning_effort="high",
|
||||
response_format=openai_response_format_schema() if response_template == "json" else None,
|
||||
safety_identifier="olmocr-bench-runner",
|
||||
)
|
||||
|
||||
# Accumulate token counts from the response
|
||||
if response.usage:
|
||||
TOTAL_INPUT_TOKENS += response.usage.prompt_tokens
|
||||
TOTAL_OUTPUT_TOKENS += response.usage.completion_tokens
|
||||
|
||||
# Increment document counter
|
||||
TOTAL_DOCUMENTS += 1
|
||||
|
||||
raw_response = response.choices[0].message.content
|
||||
|
||||
assert len(response.choices) > 0
|
||||
assert response.choices[0].message.refusal is None
|
||||
assert response.choices[0].finish_reason == "stop"
|
||||
|
||||
if response_template == "json":
|
||||
data = json.loads(raw_response)
|
||||
data = PageResponse(**data)
|
||||
|
||||
# Log token counts before returning
|
||||
logger.warning(f"Token Usage - Documents: {TOTAL_DOCUMENTS}, Input: {TOTAL_INPUT_TOKENS}, Output: {TOTAL_OUTPUT_TOKENS}")
|
||||
return data.natural_text
|
||||
else:
|
||||
# Log token counts before returning
|
||||
logger.warning(f"Token Usage - Documents: {TOTAL_DOCUMENTS}, Input: {TOTAL_INPUT_TOKENS}, Output: {TOTAL_OUTPUT_TOKENS}")
|
||||
return raw_response
|
||||
@@ -0,0 +1,61 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
from anthropic import Anthropic
|
||||
from prompts import build_openai_silver_data_prompt, claude_response_format_schema
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.prompts.anchor import get_anchor_text
|
||||
|
||||
|
||||
def run_claude(pdf_path: str, page_num: int = 1, model: str = "claude-3-7-sonnet-20250219", temperature: float = 0.1) -> str:
|
||||
"""
|
||||
Convert page of a PDF file to markdown using Claude OCR.
|
||||
This function renders the specified page of the PDF to an image, runs OCR on that image,
|
||||
and returns the OCR result as a markdown-formatted string.
|
||||
|
||||
Args:
|
||||
pdf_path (str): The local path to the PDF file.
|
||||
page_num (int): The page number to process (starting from 1).
|
||||
model (str): The Claude model to use.
|
||||
temperature (float): The temperature parameter for generation.
|
||||
|
||||
Returns:
|
||||
str: The OCR result in markdown format.
|
||||
"""
|
||||
|
||||
if not os.getenv("ANTHROPIC_API_KEY"):
|
||||
raise SystemExit("You must specify an ANTHROPIC_API_KEY")
|
||||
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=2048)
|
||||
anchor_text = get_anchor_text(pdf_path, page_num, pdf_engine="pdfreport")
|
||||
client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
|
||||
response = client.messages.create(
|
||||
model=model,
|
||||
max_tokens=3000,
|
||||
temperature=temperature,
|
||||
# system=system_prompt,
|
||||
tools=claude_response_format_schema(),
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": image_base64}},
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"{build_openai_silver_data_prompt(anchor_text)}. Use the page_response tool to respond. If the propeties are true, then extract the text from them and respond in natural_text.",
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
json_sentiment = None
|
||||
for content in response.content:
|
||||
if content.type == "tool_use" and content.name == "page_response":
|
||||
json_sentiment = content.input
|
||||
break
|
||||
|
||||
if json_sentiment:
|
||||
response = json.dumps(json_sentiment, indent=2)
|
||||
return response
|
||||
@@ -0,0 +1,76 @@
|
||||
import asyncio
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Literal
|
||||
|
||||
from pypdf import PdfReader, PdfWriter
|
||||
|
||||
|
||||
async def run_docling(
|
||||
pdf_path: str,
|
||||
page_num: int = 1,
|
||||
output_format: Literal["markdown"] = "markdown",
|
||||
use_smoldocling: bool = False,
|
||||
) -> str:
|
||||
"""Run docling CLI on a PDF file and return the results.
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: Page number to process (1-indexed)
|
||||
output_format: Output format (only markdown is supported for CLI version)
|
||||
|
||||
Returns:
|
||||
String containing the markdown output
|
||||
"""
|
||||
if output_format != "markdown":
|
||||
raise ValueError("Only markdown output format is supported for CLI version")
|
||||
|
||||
# Extract the specific page using pypdf
|
||||
pdf_reader = PdfReader(pdf_path)
|
||||
pdf_writer = PdfWriter()
|
||||
|
||||
# Convert from 1-indexed to 0-indexed
|
||||
zero_based_page_num = page_num - 1
|
||||
|
||||
if zero_based_page_num >= len(pdf_reader.pages) or zero_based_page_num < 0:
|
||||
raise ValueError(f"Page number {page_num} is out of bounds for PDF with {len(pdf_reader.pages)} pages")
|
||||
|
||||
# Add the selected page to the writer
|
||||
pdf_writer.add_page(pdf_reader.pages[zero_based_page_num])
|
||||
|
||||
# Create temporary files for the single-page PDF and output markdown
|
||||
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp_pdf_file, tempfile.NamedTemporaryFile(suffix=".md", delete=False) as tmp_md_file:
|
||||
tmp_pdf_path = tmp_pdf_file.name
|
||||
tmp_md_path = tmp_md_file.name
|
||||
|
||||
try:
|
||||
# Write the single-page PDF to the temporary file
|
||||
with open(tmp_pdf_path, "wb") as f:
|
||||
pdf_writer.write(f)
|
||||
|
||||
# Build the command to run docling on the single-page PDF
|
||||
if use_smoldocling:
|
||||
cmd = ["docling", tmp_pdf_path, "-o", tmp_md_path] # Output file
|
||||
else:
|
||||
cmd = ["docling", "--pipeline", "vlm", "--vlm-model", "smoldocling", tmp_pdf_path, "-o", tmp_md_path] # Output file
|
||||
|
||||
# Run the command asynchronously
|
||||
proc = await asyncio.create_subprocess_exec(*cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
|
||||
|
||||
stdout, stderr = await proc.communicate()
|
||||
|
||||
if proc.returncode != 0:
|
||||
error_msg = stderr.decode() if stderr else "Unknown error"
|
||||
raise RuntimeError(f"docling command failed with return code {proc.returncode}: {error_msg}")
|
||||
|
||||
# Read the results from the temporary markdown file
|
||||
with open(tmp_md_path, "r", encoding="utf-8") as f:
|
||||
result = f.read()
|
||||
|
||||
return result
|
||||
|
||||
finally:
|
||||
# Clean up the temporary files
|
||||
for path in [tmp_pdf_path, tmp_md_path]:
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
@@ -0,0 +1,190 @@
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from pypdf import PdfReader, PdfWriter
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
DOTS_OCR_PATH = REPO_ROOT / "dots.ocr"
|
||||
|
||||
if not DOTS_OCR_PATH.exists():
|
||||
raise ImportError(f"Could not find dots.ocr checkout at {DOTS_OCR_PATH}")
|
||||
|
||||
if str(DOTS_OCR_PATH) not in sys.path:
|
||||
sys.path.insert(0, str(DOTS_OCR_PATH))
|
||||
|
||||
from dots_ocr.parser import DotsOCRParser # noqa: E402
|
||||
|
||||
_PARSER_CACHE: Dict[
|
||||
Tuple[str, str, int, str, float, float, int, int, Optional[int], Optional[int], bool],
|
||||
DotsOCRParser,
|
||||
] = {}
|
||||
|
||||
|
||||
def _parse_server(server: Optional[str]) -> Tuple[str, str, int]:
|
||||
normalized = (server or "").strip() or "http://localhost:8000"
|
||||
if "://" not in normalized:
|
||||
normalized = f"http://{normalized}"
|
||||
|
||||
parsed = urlparse(normalized)
|
||||
if not parsed.scheme or not parsed.hostname:
|
||||
raise ValueError(f"Invalid server URL: {server}")
|
||||
|
||||
port = parsed.port
|
||||
if port is None:
|
||||
port = 443 if parsed.scheme == "https" else 80
|
||||
|
||||
return parsed.scheme, parsed.hostname, port
|
||||
|
||||
|
||||
def _get_parser(
|
||||
*,
|
||||
protocol: str,
|
||||
ip: str,
|
||||
port: int,
|
||||
model_name: str,
|
||||
temperature: float,
|
||||
top_p: float,
|
||||
max_completion_tokens: int,
|
||||
num_thread: int,
|
||||
dpi: int,
|
||||
min_pixels: Optional[int],
|
||||
max_pixels: Optional[int],
|
||||
use_hf: bool,
|
||||
) -> DotsOCRParser:
|
||||
cache_key = (
|
||||
protocol,
|
||||
ip,
|
||||
port,
|
||||
model_name,
|
||||
temperature,
|
||||
top_p,
|
||||
max_completion_tokens,
|
||||
num_thread,
|
||||
dpi,
|
||||
min_pixels,
|
||||
max_pixels,
|
||||
use_hf,
|
||||
)
|
||||
|
||||
parser = _PARSER_CACHE.get(cache_key)
|
||||
if parser is None:
|
||||
parser = DotsOCRParser(
|
||||
protocol=protocol,
|
||||
ip=ip,
|
||||
port=port,
|
||||
model_name=model_name,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
num_thread=num_thread,
|
||||
dpi=dpi,
|
||||
output_dir=tempfile.mkdtemp(prefix="dotsocr-cache-"),
|
||||
min_pixels=min_pixels,
|
||||
max_pixels=max_pixels,
|
||||
use_hf=use_hf,
|
||||
)
|
||||
_PARSER_CACHE[cache_key] = parser
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def _extract_markdown_from_pdf(
|
||||
parser: DotsOCRParser,
|
||||
pdf_path: str,
|
||||
page_num: int,
|
||||
prompt_mode: str,
|
||||
fitz_preprocess: bool,
|
||||
) -> str:
|
||||
if page_num < 1:
|
||||
raise ValueError("page_num must be >= 1")
|
||||
|
||||
with open(pdf_path, "rb") as pdf_file:
|
||||
reader = PdfReader(pdf_file)
|
||||
zero_index = page_num - 1
|
||||
if zero_index >= len(reader.pages):
|
||||
raise ValueError(f"Page {page_num} does not exist in {pdf_path}")
|
||||
|
||||
writer = PdfWriter()
|
||||
writer.add_page(reader.pages[zero_index])
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="dotsocr-") as tmp_dir:
|
||||
single_page_pdf = os.path.join(tmp_dir, "page.pdf")
|
||||
with open(single_page_pdf, "wb") as tmp_pdf:
|
||||
writer.write(tmp_pdf)
|
||||
|
||||
output_dir = os.path.join(tmp_dir, "output")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
results = parser.parse_file(
|
||||
single_page_pdf,
|
||||
output_dir=output_dir,
|
||||
prompt_mode=prompt_mode,
|
||||
fitz_preprocess=fitz_preprocess,
|
||||
)
|
||||
|
||||
if not results:
|
||||
raise RuntimeError("DotsOCR did not return any results")
|
||||
|
||||
# Single-page PDF means the first result corresponds to our request.
|
||||
page_result = results[0]
|
||||
md_path = page_result.get("md_content_path") or page_result.get("md_content_nohf_path")
|
||||
|
||||
if not md_path:
|
||||
raise RuntimeError(f"DotsOCR (prompt_mode={prompt_mode}) did not produce markdown output for page {page_num}")
|
||||
|
||||
if not os.path.exists(md_path):
|
||||
raise RuntimeError(f"DotsOCR reported markdown at '{md_path}', but the file does not exist")
|
||||
|
||||
with open(md_path, "r", encoding="utf-8") as md_file:
|
||||
return md_file.read()
|
||||
|
||||
|
||||
async def run_dotsocr(
|
||||
pdf_path: str,
|
||||
page_num: int = 1,
|
||||
server: Optional[str] = "http://localhost:8000",
|
||||
prompt_mode: str = "prompt_layout_all_en",
|
||||
model_name: str = "rednote-hilab/dots.ocr",
|
||||
temperature: float = 0.1,
|
||||
top_p: float = 1.0,
|
||||
max_completion_tokens: int = 16384,
|
||||
num_thread: int = 16,
|
||||
dpi: int = 200,
|
||||
min_pixels: Optional[int] = None,
|
||||
max_pixels: Optional[int] = None,
|
||||
use_hf: bool = False,
|
||||
fitz_preprocess: bool = False,
|
||||
) -> str:
|
||||
"""
|
||||
Run DotsOCR on a single PDF page and return the markdown output.
|
||||
"""
|
||||
|
||||
protocol, ip, port = _parse_server(server)
|
||||
parser = _get_parser(
|
||||
protocol=protocol,
|
||||
ip=ip,
|
||||
port=port,
|
||||
model_name=model_name,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
num_thread=num_thread,
|
||||
dpi=dpi,
|
||||
min_pixels=min_pixels,
|
||||
max_pixels=max_pixels,
|
||||
use_hf=use_hf,
|
||||
)
|
||||
|
||||
return await asyncio.to_thread(
|
||||
_extract_markdown_from_pdf,
|
||||
parser,
|
||||
pdf_path,
|
||||
page_num,
|
||||
prompt_mode,
|
||||
fitz_preprocess,
|
||||
)
|
||||
@@ -0,0 +1,163 @@
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
|
||||
from olmocr.bench.prompts import (
|
||||
build_openai_silver_data_prompt_no_document_anchoring,
|
||||
)
|
||||
from olmocr.data.renderpdf import (
|
||||
get_png_dimensions_from_base64,
|
||||
render_pdf_to_base64png,
|
||||
)
|
||||
from olmocr.prompts.anchor import get_anchor_text
|
||||
from olmocr.prompts.prompts import (
|
||||
build_openai_silver_data_prompt,
|
||||
build_openai_silver_data_prompt_v3_simple,
|
||||
)
|
||||
|
||||
# Set up logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Global variables to track token usage and documents
|
||||
TOTAL_INPUT_TOKENS = 0
|
||||
TOTAL_OUTPUT_TOKENS = 0
|
||||
TOTAL_DOCUMENTS = 0
|
||||
|
||||
|
||||
def run_gemini(
|
||||
pdf_path: str,
|
||||
page_num: int = 1,
|
||||
model: str = "gemini-2.0-flash",
|
||||
temperature: float = 0.1,
|
||||
target_longest_image_dim: int = 2048,
|
||||
prompt_template: Literal["full", "full_no_document_anchoring", "basic", "finetune", "fullv3simple"] = "finetune",
|
||||
response_template: Literal["plain", "json"] = "json",
|
||||
) -> str:
|
||||
"""
|
||||
Convert page of a PDF file to markdown using Gemini's vision capabilities.
|
||||
This function renders the specified page of the PDF to an image, runs OCR on that image,
|
||||
and returns the OCR result as a markdown-formatted string.
|
||||
|
||||
Args:
|
||||
pdf_path (str): The local path to the PDF file.
|
||||
page_num (int): The page number to process (starting from 1).
|
||||
model (str): The Gemini model to use.
|
||||
temperature (float): The temperature parameter for generation.
|
||||
|
||||
Returns:
|
||||
str: The OCR result in markdown format.
|
||||
"""
|
||||
global TOTAL_INPUT_TOKENS, TOTAL_OUTPUT_TOKENS, TOTAL_DOCUMENTS
|
||||
TOTAL_DOCUMENTS += 1
|
||||
if not os.getenv("GEMINI_API_KEY"):
|
||||
raise SystemExit("You must specify an GEMINI_API_KEY")
|
||||
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=2048)
|
||||
anchor_text = get_anchor_text(pdf_path, page_num, pdf_engine="pdfreport")
|
||||
api_key = os.getenv("GEMINI_API_KEY")
|
||||
client = genai.Client(api_key=api_key)
|
||||
image_part = types.Part(inline_data=types.Blob(mime_type="image/png", data=base64.b64decode(image_base64)))
|
||||
|
||||
if prompt_template == "full":
|
||||
text_part = types.Part(text=f"""{build_openai_silver_data_prompt(anchor_text)}""")
|
||||
elif prompt_template == "full_no_document_anchoring":
|
||||
text_part = types.Part(text=f"""{build_openai_silver_data_prompt_no_document_anchoring(anchor_text)}""")
|
||||
elif prompt_template == "fullv3simple":
|
||||
width, height = get_png_dimensions_from_base64(image_base64)
|
||||
prompt = build_openai_silver_data_prompt_v3_simple(width, height)
|
||||
text_part = types.Part(text=prompt)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
if response_template == "json":
|
||||
generation_config = types.GenerateContentConfig(
|
||||
temperature=temperature,
|
||||
top_p=1.0,
|
||||
top_k=32,
|
||||
max_output_tokens=10000,
|
||||
response_mime_type="application/json",
|
||||
response_schema=genai.types.Schema(
|
||||
type=genai.types.Type.OBJECT,
|
||||
required=["primary_language", "is_rotation_valid", "rotation_correction", "is_table", "is_diagram", "natural_text"],
|
||||
properties={
|
||||
"primary_language": genai.types.Schema(
|
||||
type=genai.types.Type.STRING,
|
||||
),
|
||||
"is_rotation_valid": genai.types.Schema(
|
||||
type=genai.types.Type.BOOLEAN,
|
||||
),
|
||||
"rotation_correction": genai.types.Schema(
|
||||
type=genai.types.Type.STRING,
|
||||
enum=["0", "90", "180", "270"],
|
||||
),
|
||||
"is_table": genai.types.Schema(
|
||||
type=genai.types.Type.BOOLEAN,
|
||||
),
|
||||
"is_diagram": genai.types.Schema(
|
||||
type=genai.types.Type.BOOLEAN,
|
||||
),
|
||||
"natural_text": genai.types.Schema(
|
||||
type=genai.types.Type.STRING,
|
||||
),
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
response = client.models.generate_content(
|
||||
model=f"models/{model}",
|
||||
contents=[types.Content(parts=[image_part, text_part])],
|
||||
config=generation_config,
|
||||
)
|
||||
|
||||
assert len(response.candidates) > 0, "No candidates found"
|
||||
assert response.candidates[0].finish_reason == types.FinishReason.STOP, "Finish reason was not STOP, likely a processing error or repetition failure"
|
||||
|
||||
# Extract token counts from usage metadata
|
||||
if hasattr(response, "usage_metadata"):
|
||||
input_tokens = getattr(response.usage_metadata, "prompt_token_count", 0)
|
||||
output_tokens = getattr(response.usage_metadata, "candidates_token_count", 0)
|
||||
TOTAL_INPUT_TOKENS += input_tokens
|
||||
TOTAL_OUTPUT_TOKENS += output_tokens
|
||||
|
||||
result = response.candidates[0].content.parts[0].text
|
||||
parsed = json.loads(result)
|
||||
|
||||
# The json schema is slightly off with gemini vs chatgpt, so we don't verify it
|
||||
logger.warning(
|
||||
f"[Before Return - JSON] Total Documents: {TOTAL_DOCUMENTS}, Total Input Tokens: {TOTAL_INPUT_TOKENS}, Total Output Tokens: {TOTAL_OUTPUT_TOKENS}"
|
||||
)
|
||||
return parsed["natural_text"]
|
||||
else:
|
||||
generation_config = types.GenerateContentConfig(
|
||||
temperature=temperature,
|
||||
top_p=1.0,
|
||||
top_k=32,
|
||||
max_output_tokens=4096,
|
||||
)
|
||||
|
||||
response = client.models.generate_content(
|
||||
model=f"models/{model}",
|
||||
contents=[types.Content(parts=[image_part, text_part])],
|
||||
config=generation_config,
|
||||
)
|
||||
|
||||
assert len(response.candidates) > 0, "No candidates found"
|
||||
assert response.candidates[0].finish_reason == types.FinishReason.STOP, "Finish reason was not STOP, likely a processing error or repetition failure"
|
||||
|
||||
# Extract token counts from usage metadata
|
||||
if hasattr(response, "usage_metadata"):
|
||||
input_tokens = getattr(response.usage_metadata, "prompt_token_count", 0)
|
||||
output_tokens = getattr(response.usage_metadata, "candidates_token_count", 0)
|
||||
TOTAL_INPUT_TOKENS += input_tokens
|
||||
TOTAL_OUTPUT_TOKENS += output_tokens
|
||||
|
||||
result = response.candidates[0].content.parts[0].text
|
||||
logger.warning(
|
||||
f"[Before Return - Plain] Total Documents: {TOTAL_DOCUMENTS}, Total Input Tokens: {TOTAL_INPUT_TOKENS}, Total Output Tokens: {TOTAL_OUTPUT_TOKENS}"
|
||||
)
|
||||
return result
|
||||
@@ -0,0 +1,67 @@
|
||||
import base64
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
|
||||
# Global cache for the model and tokenizer.
|
||||
_device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
_model = None
|
||||
_tokenizer = None
|
||||
|
||||
|
||||
def load_model():
|
||||
"""
|
||||
Load the GOT-OCR model and tokenizer if they haven't been loaded already.
|
||||
Returns:
|
||||
model: The GOT-OCR model loaded on the appropriate device.
|
||||
tokenizer: The corresponding tokenizer.
|
||||
"""
|
||||
global _model, _tokenizer
|
||||
if _model is None or _tokenizer is None:
|
||||
_tokenizer = AutoTokenizer.from_pretrained("ucaslcl/GOT-OCR2_0", trust_remote_code=True)
|
||||
_model = AutoModel.from_pretrained(
|
||||
"ucaslcl/GOT-OCR2_0",
|
||||
trust_remote_code=True,
|
||||
use_safetensors=True,
|
||||
revision="979938bf89ccdc949c0131ddd3841e24578a4742",
|
||||
pad_token_id=_tokenizer.eos_token_id,
|
||||
)
|
||||
_model = _model.eval().to(_device)
|
||||
return _model, _tokenizer
|
||||
|
||||
|
||||
def run_gotocr(pdf_path: str, page_num: int = 1, ocr_type: str = "ocr") -> str:
|
||||
"""
|
||||
Convert page of a PDF file to markdown using GOT-OCR.
|
||||
|
||||
This function renders the first page of the PDF to an image, runs OCR on that image,
|
||||
and returns the OCR result as a markdown-formatted string.
|
||||
|
||||
Args:
|
||||
pdf_path (str): The local path to the PDF file.
|
||||
|
||||
Returns:
|
||||
str: The OCR result in markdown format.
|
||||
"""
|
||||
# Ensure the model is loaded (cached across calls)
|
||||
model, tokenizer = load_model()
|
||||
|
||||
# Convert the first page of the PDF to a base64-encoded PNG image.
|
||||
base64image = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=1024)
|
||||
|
||||
# Write the image to a temporary file.
|
||||
with tempfile.NamedTemporaryFile("wb", suffix=".png", delete=False) as tmp:
|
||||
tmp.write(base64.b64decode(base64image))
|
||||
tmp_filename = tmp.name
|
||||
|
||||
# Run GOT-OCR on the saved image.
|
||||
result = model.chat(tokenizer, tmp_filename, ocr_type=ocr_type)
|
||||
|
||||
# Clean up the temporary file.
|
||||
os.remove(tmp_filename)
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,69 @@
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
from marker.config.parser import ConfigParser
|
||||
from marker.converters.pdf import PdfConverter
|
||||
from marker.models import create_model_dict
|
||||
from marker.output import text_from_rendered
|
||||
from pypdf import PdfReader, PdfWriter
|
||||
|
||||
_marker_converter = None
|
||||
|
||||
|
||||
def run_marker(pdf_path: str, page_num: int = 1) -> str:
|
||||
global _marker_converter
|
||||
|
||||
if _marker_converter is None:
|
||||
# Create a configuration dictionary with the necessary settings
|
||||
config = {
|
||||
"force_ocr": True, # This enables conversion of inline math to LaTeX
|
||||
"use_llm": False, # We would prefer to run just plain marker for reporting bench results, not hybrid mode
|
||||
"disable_tqdm": True, # Disable tqdm for cleaner output
|
||||
"recognition_batch_size": 256,
|
||||
"layout_batch_size": 48,
|
||||
"detection_batch_size": 48,
|
||||
"equation_batch_size": 64,
|
||||
"table_rec_batch_size": 48,
|
||||
"ocr_error_batch_size": 64,
|
||||
}
|
||||
config_parser = ConfigParser(config)
|
||||
|
||||
_marker_converter = PdfConverter(
|
||||
artifact_dict=create_model_dict(),
|
||||
config=config_parser.generate_config_dict(),
|
||||
)
|
||||
|
||||
# Extract the specific page from the PDF
|
||||
pdf_to_process = pdf_path
|
||||
temp_file = None
|
||||
|
||||
if page_num > 0: # If a specific page is requested
|
||||
reader = PdfReader(pdf_path)
|
||||
|
||||
# Check if the requested page exists
|
||||
if page_num > len(reader.pages):
|
||||
raise ValueError(f"Page {page_num} does not exist in the PDF. PDF has {len(reader.pages)} pages.")
|
||||
|
||||
# Create a new PDF with just the requested page
|
||||
writer = PdfWriter()
|
||||
# pypdf uses 0-based indexing, so subtract 1 from page_num
|
||||
writer.add_page(reader.pages[page_num - 1])
|
||||
|
||||
# Save the extracted page to a temporary file
|
||||
temp_file = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False)
|
||||
temp_file.close() # Close the file but keep the name
|
||||
|
||||
with open(temp_file.name, "wb") as output_pdf:
|
||||
writer.write(output_pdf)
|
||||
|
||||
pdf_to_process = temp_file.name
|
||||
|
||||
try:
|
||||
# Process the PDF (either original or single-page extract)
|
||||
rendered = _marker_converter(pdf_to_process)
|
||||
text, _, images = text_from_rendered(rendered)
|
||||
return text
|
||||
finally:
|
||||
# Clean up the temporary file if it was created
|
||||
if temp_file and os.path.exists(temp_file.name):
|
||||
os.unlink(temp_file.name)
|
||||
@@ -0,0 +1,74 @@
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
from magic_pdf.config.enums import SupportedPdfParseMethod
|
||||
from magic_pdf.data.data_reader_writer import FileBasedDataReader, FileBasedDataWriter
|
||||
from magic_pdf.data.dataset import PymuDocDataset
|
||||
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
|
||||
from pypdf import PdfReader, PdfWriter
|
||||
|
||||
|
||||
def run_mineru(pdf_path: str, page_num: int = 1) -> str:
|
||||
output_folder = tempfile.TemporaryDirectory()
|
||||
image_output_folder = tempfile.TemporaryDirectory()
|
||||
|
||||
# Initialize writers (same for all PDFs)
|
||||
image_writer = FileBasedDataWriter(image_output_folder.name)
|
||||
md_writer = FileBasedDataWriter(output_folder.name)
|
||||
|
||||
if page_num > 0: # If a specific page is requested
|
||||
reader = PdfReader(pdf_path)
|
||||
|
||||
# Check if the requested page exists
|
||||
if page_num > len(reader.pages):
|
||||
raise ValueError(f"Page {page_num} does not exist in the PDF. PDF has {len(reader.pages)} pages.")
|
||||
|
||||
# Create a new PDF with just the requested page
|
||||
writer = PdfWriter()
|
||||
# pypdf uses 0-based indexing, so subtract 1 from page_num
|
||||
writer.add_page(reader.pages[page_num - 1])
|
||||
|
||||
# Save the extracted page to a temporary file
|
||||
temp_file = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False)
|
||||
temp_file.close() # Close the file but keep the name
|
||||
|
||||
with open(temp_file.name, "wb") as output_pdf:
|
||||
writer.write(output_pdf)
|
||||
|
||||
pdf_to_process = temp_file.name
|
||||
else:
|
||||
pdf_to_process = pdf_path
|
||||
|
||||
try:
|
||||
# Read the PDF file bytes
|
||||
reader = FileBasedDataReader("")
|
||||
pdf_bytes = reader.read(pdf_to_process)
|
||||
|
||||
# Create dataset instance
|
||||
ds = PymuDocDataset(pdf_bytes)
|
||||
|
||||
# Inference: decide whether to run OCR mode based on dataset classification
|
||||
if ds.classify() == SupportedPdfParseMethod.OCR:
|
||||
infer_result = ds.apply(doc_analyze, ocr=True)
|
||||
pipe_result = infer_result.pipe_ocr_mode(image_writer)
|
||||
else:
|
||||
infer_result = ds.apply(doc_analyze, ocr=False)
|
||||
pipe_result = infer_result.pipe_txt_mode(image_writer)
|
||||
|
||||
# Generate markdown content; the image directory is the basename of the images output folder
|
||||
image_dir_basename = os.path.basename(image_output_folder.name)
|
||||
# md_content = pipe_result.get_markdown(image_dir_basename)
|
||||
|
||||
# Dump markdown file
|
||||
with tempfile.NamedTemporaryFile("w+", suffix="md") as tf:
|
||||
pipe_result.dump_md(md_writer, tf.name, image_dir_basename)
|
||||
tf.flush()
|
||||
|
||||
tf.seek(0)
|
||||
md_data = tf.read()
|
||||
|
||||
return md_data
|
||||
finally:
|
||||
# Clean up the temporary file if it was created
|
||||
if temp_file and os.path.exists(temp_file.name):
|
||||
os.unlink(temp_file.name)
|
||||
@@ -0,0 +1,74 @@
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
from mistralai import Mistral
|
||||
from pypdf import PdfReader, PdfWriter
|
||||
|
||||
|
||||
def run_mistral(pdf_path: str, page_num: int = 1) -> str:
|
||||
"""
|
||||
Convert page of a PDF file to markdown using the mistral OCR api
|
||||
https://docs.mistral.ai/capabilities/document/
|
||||
|
||||
Args:
|
||||
pdf_path (str): The local path to the PDF file.
|
||||
|
||||
Returns:
|
||||
str: The OCR result in markdown format.
|
||||
"""
|
||||
if not os.getenv("MISTRAL_API_KEY"):
|
||||
raise SystemExit("You must specify an MISTRAL_API_KEY")
|
||||
|
||||
api_key = os.environ["MISTRAL_API_KEY"]
|
||||
client = Mistral(api_key=api_key)
|
||||
|
||||
if page_num > 0: # If a specific page is requested
|
||||
reader = PdfReader(pdf_path)
|
||||
|
||||
# Check if the requested page exists
|
||||
if page_num > len(reader.pages):
|
||||
raise ValueError(f"Page {page_num} does not exist in the PDF. PDF has {len(reader.pages)} pages.")
|
||||
|
||||
# Create a new PDF with just the requested page
|
||||
writer = PdfWriter()
|
||||
# pypdf uses 0-based indexing, so subtract 1 from page_num
|
||||
writer.add_page(reader.pages[page_num - 1])
|
||||
|
||||
# Save the extracted page to a temporary file
|
||||
temp_file = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False)
|
||||
temp_file.close() # Close the file but keep the name
|
||||
|
||||
with open(temp_file.name, "wb") as output_pdf:
|
||||
writer.write(output_pdf)
|
||||
|
||||
pdf_to_process = temp_file.name
|
||||
else:
|
||||
pdf_to_process = pdf_path
|
||||
|
||||
try:
|
||||
with open(pdf_to_process, "rb") as pf:
|
||||
uploaded_pdf = client.files.upload(
|
||||
file={
|
||||
"file_name": os.path.basename(pdf_path),
|
||||
"content": pf,
|
||||
},
|
||||
purpose="ocr",
|
||||
)
|
||||
|
||||
signed_url = client.files.get_signed_url(file_id=uploaded_pdf.id)
|
||||
|
||||
ocr_response = client.ocr.process(
|
||||
model="mistral-ocr-2503",
|
||||
document={
|
||||
"type": "document_url",
|
||||
"document_url": signed_url.url,
|
||||
},
|
||||
)
|
||||
|
||||
client.files.delete(file_id=uploaded_pdf.id)
|
||||
|
||||
return ocr_response.pages[0].markdown
|
||||
finally:
|
||||
# Clean up the temporary file if it was created
|
||||
if temp_file and os.path.exists(temp_file.name):
|
||||
os.unlink(temp_file.name)
|
||||
@@ -0,0 +1,89 @@
|
||||
import base64
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import AutoModelForImageTextToText, AutoProcessor, AutoTokenizer
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
|
||||
_model = None
|
||||
_tokenizer = None
|
||||
_processor = None
|
||||
_device = None
|
||||
|
||||
|
||||
def load_model(model_path: str = "nanonets/Nanonets-OCR-s"):
|
||||
global _model, _tokenizer, _processor, _device
|
||||
|
||||
if _model is None:
|
||||
_device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
_model = AutoModelForImageTextToText.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype="auto",
|
||||
device_map="auto",
|
||||
# attn_implementation="flash_attention_2"
|
||||
)
|
||||
_model.eval()
|
||||
_tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
_processor = AutoProcessor.from_pretrained(model_path)
|
||||
|
||||
return _model, _tokenizer, _processor
|
||||
|
||||
|
||||
async def run_nanonetsocr(pdf_path: str, page_num: int = 1, model_path: str = "nanonets/Nanonets-OCR-s", max_new_tokens: int = 4096, **kwargs) -> str:
|
||||
"""
|
||||
Convert page of a PDF file to markdown using NANONETS-OCR.
|
||||
|
||||
This function renders the first page of the PDF to an image, runs OCR on that image,
|
||||
and returns the OCR result as a markdown-formatted string.
|
||||
|
||||
Args:
|
||||
pdf_path (str): The local path to the PDF file.
|
||||
|
||||
Returns:
|
||||
str: The OCR result in markdown format.
|
||||
"""
|
||||
|
||||
model, tokenizer, processor = load_model(model_path)
|
||||
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=1024)
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
|
||||
image_data = base64.b64decode(image_base64)
|
||||
temp_file.write(image_data)
|
||||
temp_image_path = temp_file.name
|
||||
|
||||
try:
|
||||
image = Image.open(temp_image_path)
|
||||
prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "image": f"file://{temp_image_path}"},
|
||||
{"type": "text", "text": prompt},
|
||||
],
|
||||
},
|
||||
]
|
||||
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||
|
||||
inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt", use_fast=True)
|
||||
inputs = inputs.to(model.device)
|
||||
with torch.no_grad():
|
||||
output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
|
||||
|
||||
generated_ids = [output_ids[len(input_ids) :] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
|
||||
output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
|
||||
cleaned_text = re.sub(r"<page_number>\d+</page_number>", "", output_text[0])
|
||||
|
||||
return cleaned_text
|
||||
|
||||
finally:
|
||||
try:
|
||||
os.unlink(temp_image_path)
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to remove temporary file {temp_image_path}: {e}")
|
||||
@@ -0,0 +1,86 @@
|
||||
import base64
|
||||
import subprocess
|
||||
|
||||
import httpx
|
||||
|
||||
from olmocr.data.renderpdf import get_pdf_media_box_width_height
|
||||
|
||||
|
||||
# Logic to set min size from here: https://github.com/NanoNets/Nanonets-OCR2/blob/main/Nanonets-OCR2-Cookbook/image2md.ipynb
|
||||
def render_pdf_to_base64png_min_short_size(local_pdf_path: str, page_num: int, target_shortest_dim: int = 2048) -> str:
|
||||
shortest_dim = min(get_pdf_media_box_width_height(local_pdf_path, page_num))
|
||||
|
||||
# Convert PDF page to PNG using pdftoppm
|
||||
pdftoppm_result = subprocess.run(
|
||||
[
|
||||
"pdftoppm",
|
||||
"-png",
|
||||
"-f",
|
||||
str(page_num),
|
||||
"-l",
|
||||
str(page_num),
|
||||
"-r",
|
||||
str(target_shortest_dim * 72 / shortest_dim), # 72 pixels per point is the conversion factor
|
||||
local_pdf_path,
|
||||
],
|
||||
timeout=120,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
assert pdftoppm_result.returncode == 0, pdftoppm_result.stderr
|
||||
return base64.b64encode(pdftoppm_result.stdout).decode("utf-8")
|
||||
|
||||
|
||||
async def run_server(
|
||||
pdf_path: str,
|
||||
page_num: int = 1,
|
||||
server: str = "localhost:30000",
|
||||
model: str = "nanonets/Nanonets-OCR2-3B",
|
||||
temperature: float = 0.0,
|
||||
page_dimensions: int = 1280,
|
||||
) -> str:
|
||||
"""
|
||||
Convert page of a PDF file to markdown by calling a request
|
||||
running against an openai compatible server.
|
||||
|
||||
You can use this for running against vllm, sglang, servers
|
||||
as well as mixing and matching different model's.
|
||||
|
||||
It will only make one direct request, with no retries or error checking.
|
||||
|
||||
Returns:
|
||||
str: The OCR result in markdown format.
|
||||
"""
|
||||
# Convert the first page of the PDF to a base64-encoded PNG image.
|
||||
image_base64 = render_pdf_to_base64png_min_short_size(pdf_path, page_num=page_num, target_shortest_dim=page_dimensions)
|
||||
|
||||
# Now use th
|
||||
prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""
|
||||
|
||||
request = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
|
||||
{"type": "text", "text": prompt},
|
||||
],
|
||||
},
|
||||
],
|
||||
"temperature": temperature,
|
||||
"max_tokens": 4096,
|
||||
}
|
||||
|
||||
# Make request and get response using httpx
|
||||
url = f"http://{server}/v1/chat/completions"
|
||||
|
||||
async with httpx.AsyncClient(timeout=300) as client:
|
||||
response = await client.post(url, json=request)
|
||||
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
choice = data["choices"][0]
|
||||
return choice["message"]["content"]
|
||||
@@ -0,0 +1,112 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
# Import necessary components from olmocr
|
||||
from olmocr.pipeline import (
|
||||
MetricsKeeper,
|
||||
PageResult,
|
||||
WorkerTracker,
|
||||
process_page,
|
||||
vllm_server_host,
|
||||
vllm_server_ready,
|
||||
)
|
||||
|
||||
# Setup basic logging
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
||||
logger = logging.getLogger("olmocr_runner")
|
||||
|
||||
|
||||
# Basic configuration
|
||||
@dataclass
|
||||
class Args:
|
||||
model: str = "allenai/olmOCR-2-7B-1025-FP8"
|
||||
server: str = "http://localhost:30044/v1"
|
||||
port: int = 30044
|
||||
model_chat_template: str = "qwen2-vl"
|
||||
max_model_len: int = 16384
|
||||
guided_decoding: bool = False
|
||||
gpu_memory_utilization: float = 0.8
|
||||
target_longest_image_dim: int = 1288
|
||||
target_anchor_text_len: int = -1
|
||||
max_page_retries: int = 8
|
||||
max_page_error_rate: float = 0.004
|
||||
tensor_parallel_size: int = 1
|
||||
data_parallel_size: int = 1
|
||||
|
||||
|
||||
server_check_lock = asyncio.Lock()
|
||||
|
||||
|
||||
async def run_olmocr_pipeline(pdf_path: str, page_num: int = 1, model: str = "allenai/olmOCR-2-7B-1025-FP8") -> Optional[str]:
|
||||
"""
|
||||
Process a single page of a PDF using the official olmocr pipeline's process_page function
|
||||
|
||||
Args:
|
||||
pdf_path: Path to the PDF file
|
||||
page_num: Page number to process (1-indexed)
|
||||
|
||||
Returns:
|
||||
The extracted text from the page or None if processing failed
|
||||
"""
|
||||
# Ensure global variables are initialized
|
||||
global metrics, tracker
|
||||
if "metrics" not in globals() or metrics is None:
|
||||
metrics = MetricsKeeper(window=60 * 5)
|
||||
if "tracker" not in globals() or tracker is None:
|
||||
tracker = WorkerTracker()
|
||||
|
||||
args = Args()
|
||||
args.model = model
|
||||
semaphore = asyncio.Semaphore(1)
|
||||
worker_id = 0 # Using 0 as default worker ID
|
||||
|
||||
# Ensure server is running
|
||||
async with server_check_lock:
|
||||
_server_task = None
|
||||
try:
|
||||
await asyncio.wait_for(vllm_server_ready(args), timeout=5)
|
||||
logger.info("Using existing vllm server")
|
||||
except Exception:
|
||||
logger.info("Starting new vllm server")
|
||||
_server_task = asyncio.create_task(vllm_server_host(args.model, args, semaphore))
|
||||
await vllm_server_ready(args)
|
||||
|
||||
# Sets the model name used in the pipeline code, it's a hack sadly
|
||||
args.model = "olmocr"
|
||||
|
||||
try:
|
||||
# Process the page using the pipeline's process_page function
|
||||
# Note: process_page expects both original path and local path
|
||||
# In our case, we're using the same path for both
|
||||
page_result: PageResult = await process_page(args=args, worker_id=worker_id, pdf_orig_path=pdf_path, pdf_local_path=pdf_path, page_num=page_num)
|
||||
|
||||
# Return the natural text from the response
|
||||
if page_result and page_result.response and not page_result.is_fallback:
|
||||
return page_result.response.natural_text
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing page: {type(e).__name__} - {str(e)}")
|
||||
return None
|
||||
|
||||
finally:
|
||||
# We leave the server running for potential reuse
|
||||
pass
|
||||
|
||||
|
||||
async def main():
|
||||
# Example usage
|
||||
pdf_path = "your_pdf_path.pdf"
|
||||
page_num = 1
|
||||
|
||||
result = await run_olmocr_pipeline(pdf_path, page_num)
|
||||
if result:
|
||||
print(f"Extracted text: {result[:200]}...") # Print first 200 chars
|
||||
else:
|
||||
print("Failed to extract text from the page")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,31 @@
|
||||
from threading import Lock
|
||||
|
||||
from paddleocr import PPStructureV3
|
||||
|
||||
# Run's paddle paddle as in the docs here: https://huggingface.co/PaddlePaddle/PP-OCRv5_server_det
|
||||
# text_detection_model_name="PP-OCRv5_server_det",
|
||||
# and using the PP-StructureV3 pipeline to create markdown
|
||||
paddle_pipeline = None
|
||||
paddle_pipeline_lock = Lock()
|
||||
|
||||
|
||||
def run_paddlepaddle(pdf_path: str, page_num: int = 1, **kwargs) -> str:
|
||||
global paddle_pipeline
|
||||
|
||||
with paddle_pipeline_lock:
|
||||
if paddle_pipeline is None:
|
||||
paddle_pipeline = PPStructureV3(
|
||||
text_detection_model_name="PP-OCRv5_server_det",
|
||||
use_doc_orientation_classify=False, # Use use_doc_orientation_classify to enable/disable document orientation classification model
|
||||
use_doc_unwarping=False, # Use use_doc_unwarping to enable/disable document unwarping module
|
||||
use_textline_orientation=False, # Use use_textline_orientation to enable/disable textline orientation classification model
|
||||
device="gpu:0", # Use device to specify GPU for model inference
|
||||
)
|
||||
|
||||
output = paddle_pipeline.predict(pdf_path)
|
||||
result = ""
|
||||
for cur_page_0_indexed, res in enumerate(output):
|
||||
if cur_page_0_indexed == page_num - 1:
|
||||
result = res.markdown["markdown_texts"]
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,23 @@
|
||||
from threading import Lock
|
||||
|
||||
from paddleocr import PaddleOCRVL
|
||||
|
||||
# Using docs from here: https://huggingface.co/PaddlePaddle/PaddleOCR-VL
|
||||
paddle_pipeline = None
|
||||
paddle_pipeline_lock = Lock()
|
||||
|
||||
|
||||
def run_paddlevl(pdf_path: str, page_num: int = 1, **kwargs) -> str:
|
||||
global paddle_pipeline
|
||||
|
||||
with paddle_pipeline_lock:
|
||||
if paddle_pipeline is None:
|
||||
paddle_pipeline = PaddleOCRVL()
|
||||
|
||||
output = paddle_pipeline.predict(pdf_path)
|
||||
result = ""
|
||||
for cur_page_0_indexed, res in enumerate(output):
|
||||
if cur_page_0_indexed == page_num - 1:
|
||||
result = res.markdown["markdown_texts"]
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,58 @@
|
||||
import httpx
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
|
||||
|
||||
async def run_rolmocr(
|
||||
pdf_path: str,
|
||||
page_num: int = 1,
|
||||
server: str = "localhost:30000",
|
||||
model: str = "reducto/RolmOCR",
|
||||
temperature: float = 0.2,
|
||||
target_longest_image_dim: int = 1024,
|
||||
) -> str:
|
||||
"""
|
||||
|
||||
|
||||
Returns:
|
||||
str: The OCR result in markdown format.
|
||||
"""
|
||||
# Convert the first page of the PDF to a base64-encoded PNG image.
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=target_longest_image_dim)
|
||||
|
||||
request = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/png;base64,{image_base64}"},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Return the plain text representation of this document as if you were reading it naturally.\n",
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
"temperature": temperature,
|
||||
"max_tokens": 4096,
|
||||
}
|
||||
|
||||
# Make request and get response using httpx
|
||||
url = f"http://{server}/v1/chat/completions"
|
||||
|
||||
async with httpx.AsyncClient(timeout=300) as client:
|
||||
response = await client.post(url, json=request)
|
||||
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
choice = data["choices"][0]
|
||||
assert (
|
||||
choice["finish_reason"] == "stop"
|
||||
), "Response from server did not finish with finish_reason stop as expected, this is probably going to lead to bad data"
|
||||
|
||||
return choice["message"]["content"]
|
||||
@@ -0,0 +1,100 @@
|
||||
import json
|
||||
from typing import Literal
|
||||
|
||||
import httpx
|
||||
|
||||
from olmocr.bench.prompts import (
|
||||
build_basic_prompt,
|
||||
build_openai_silver_data_prompt_no_document_anchoring,
|
||||
)
|
||||
from olmocr.data.renderpdf import (
|
||||
get_png_dimensions_from_base64,
|
||||
render_pdf_to_base64png,
|
||||
)
|
||||
from olmocr.prompts.anchor import get_anchor_text
|
||||
from olmocr.prompts.prompts import (
|
||||
PageResponse,
|
||||
build_finetuning_prompt,
|
||||
build_no_anchoring_v4_yaml_prompt,
|
||||
build_openai_silver_data_prompt,
|
||||
build_openai_silver_data_prompt_v3_simple,
|
||||
)
|
||||
|
||||
|
||||
async def run_server(
|
||||
pdf_path: str,
|
||||
page_num: int = 1,
|
||||
endpoint: str = "http://localhost:8000/v1",
|
||||
model: str = "allenai/olmOCR-7B-0225-preview",
|
||||
temperature: float = 0.1,
|
||||
target_longest_image_dim: int = 1024,
|
||||
prompt_template: Literal["full", "full_no_document_anchoring", "fullv3simple", "finetune_v4_yaml", "basic", "finetune"] = "fullv3simple",
|
||||
response_template: Literal["plain", "json"] = "plain",
|
||||
) -> str:
|
||||
"""
|
||||
Convert page of a PDF file to markdown by calling a request
|
||||
running against an openai compatible server.
|
||||
|
||||
You can use this for running against vllm, sglang, servers
|
||||
as well as mixing and matching different model's.
|
||||
|
||||
It will only make one direct request, with no retries or error checking.
|
||||
|
||||
Returns:
|
||||
str: The OCR result in markdown format.
|
||||
"""
|
||||
# Convert the first page of the PDF to a base64-encoded PNG image.
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=target_longest_image_dim)
|
||||
anchor_text = get_anchor_text(pdf_path, page_num, pdf_engine="pdfreport")
|
||||
|
||||
if prompt_template == "full":
|
||||
prompt = build_openai_silver_data_prompt(anchor_text)
|
||||
elif prompt_template == "full_no_document_anchoring":
|
||||
prompt = build_openai_silver_data_prompt_no_document_anchoring(anchor_text)
|
||||
elif prompt_template == "finetune":
|
||||
prompt = build_finetuning_prompt(anchor_text)
|
||||
elif prompt_template == "basic":
|
||||
prompt = build_basic_prompt()
|
||||
elif prompt_template == "finetune_v4_yaml":
|
||||
prompt = build_no_anchoring_v4_yaml_prompt()
|
||||
elif prompt_template == "fullv3simple":
|
||||
width, height = get_png_dimensions_from_base64(image_base64)
|
||||
prompt = build_openai_silver_data_prompt_v3_simple(width, height)
|
||||
else:
|
||||
raise ValueError("Unknown prompt template")
|
||||
|
||||
request = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
"temperature": temperature,
|
||||
"max_tokens": 8000,
|
||||
}
|
||||
|
||||
# Make request and get response using httpx
|
||||
url = f"{endpoint.rstrip('/')}/chat/completions"
|
||||
|
||||
async with httpx.AsyncClient(timeout=300) as client:
|
||||
response = await client.post(url, json=request)
|
||||
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
choice = data["choices"][0]
|
||||
assert (
|
||||
choice["finish_reason"] == "stop"
|
||||
), "Response from server did not finish with finish_reason stop as expected, this is probably going to lead to bad data"
|
||||
|
||||
if response_template == "json":
|
||||
page_data = json.loads(choice["message"]["content"])
|
||||
page_response = PageResponse(**page_data)
|
||||
return page_response.natural_text
|
||||
elif response_template == "plain":
|
||||
return choice["message"]["content"]
|
||||
@@ -0,0 +1,129 @@
|
||||
import base64
|
||||
import json
|
||||
from io import BytesIO
|
||||
from typing import Literal
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import (
|
||||
AutoProcessor,
|
||||
Qwen2_5_VLForConditionalGeneration,
|
||||
)
|
||||
|
||||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||||
from olmocr.prompts.anchor import get_anchor_text
|
||||
from olmocr.prompts.prompts import (
|
||||
PageResponse,
|
||||
build_finetuning_prompt,
|
||||
build_no_anchoring_yaml_prompt,
|
||||
build_openai_silver_data_prompt,
|
||||
)
|
||||
from olmocr.train.front_matter import FrontMatterParser
|
||||
|
||||
_cached_model = None
|
||||
_cached_processor = None
|
||||
|
||||
|
||||
def run_transformers(
|
||||
pdf_path: str,
|
||||
page_num: int = 1,
|
||||
model_name: str = "allenai/olmOCR-7B-0725-FP8",
|
||||
temperature: float = 0.1,
|
||||
target_longest_image_dim: int = 1024,
|
||||
prompt_template: Literal["full", "finetune", "yaml"] = "yaml",
|
||||
response_template: Literal["plain", "json", "yaml"] = "yaml",
|
||||
) -> str:
|
||||
"""
|
||||
Convert page of a PDF file to markdown by calling a request
|
||||
running against an openai compatible server.
|
||||
|
||||
You can use this for running against vllm, sglang, servers
|
||||
as well as mixing and matching different model's.
|
||||
|
||||
It will only make one direct request, with no retries or error checking.
|
||||
|
||||
Returns:
|
||||
str: The OCR result in markdown format.
|
||||
"""
|
||||
# Initialize the model
|
||||
global _cached_model, _cached_processor
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
if _cached_model is None:
|
||||
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
||||
model_name, torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="flash_attention_2"
|
||||
).eval()
|
||||
processor = AutoProcessor.from_pretrained(model_name)
|
||||
|
||||
model = model.to(device)
|
||||
|
||||
_cached_model = model
|
||||
_cached_processor = processor
|
||||
else:
|
||||
model = _cached_model
|
||||
processor = _cached_processor
|
||||
|
||||
# Convert the first page of the PDF to a base64-encoded PNG image.
|
||||
image_base64 = render_pdf_to_base64png(pdf_path, page_num=page_num, target_longest_image_dim=target_longest_image_dim)
|
||||
|
||||
if prompt_template == "yaml":
|
||||
prompt = build_no_anchoring_yaml_prompt()
|
||||
else:
|
||||
anchor_text = get_anchor_text(pdf_path, page_num, pdf_engine="pdfreport")
|
||||
if prompt_template == "full":
|
||||
prompt = build_openai_silver_data_prompt(anchor_text)
|
||||
else:
|
||||
prompt = build_finetuning_prompt(anchor_text)
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
|
||||
{"type": "text", "text": prompt},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
# Apply the chat template and processor
|
||||
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
||||
main_image = Image.open(BytesIO(base64.b64decode(image_base64)))
|
||||
|
||||
inputs = processor(
|
||||
text=[text],
|
||||
images=[main_image],
|
||||
padding=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
inputs = {key: value.to(device) for (key, value) in inputs.items()}
|
||||
|
||||
# Generate the output
|
||||
MAX_NEW_TOKENS = 3000
|
||||
with torch.no_grad():
|
||||
output = model.generate(
|
||||
**inputs,
|
||||
temperature=temperature,
|
||||
max_new_tokens=MAX_NEW_TOKENS,
|
||||
num_return_sequences=1,
|
||||
do_sample=True,
|
||||
)
|
||||
|
||||
# Decode the output
|
||||
prompt_length = inputs["input_ids"].shape[1]
|
||||
new_tokens = output[:, prompt_length:]
|
||||
text_output = processor.tokenizer.batch_decode(new_tokens, skip_special_tokens=True)[0]
|
||||
|
||||
assert new_tokens.shape[1] < MAX_NEW_TOKENS, "Output exceed max new tokens"
|
||||
|
||||
if response_template == "json":
|
||||
page_data = json.loads(text_output)
|
||||
page_response = PageResponse(**page_data)
|
||||
return page_response.natural_text if page_response.natural_text else ""
|
||||
elif response_template == "yaml":
|
||||
# Parse YAML front matter and extract natural text
|
||||
parser = FrontMatterParser(front_matter_class=PageResponse)
|
||||
front_matter, text = parser._extract_front_matter_and_text(text_output)
|
||||
page_response = parser._parse_front_matter(front_matter, text)
|
||||
return page_response.natural_text if page_response.natural_text else ""
|
||||
elif response_template == "plain":
|
||||
return text_output
|
||||
@@ -0,0 +1 @@
|
||||
{"pdf": "blank_book_pg1.pdf", "page": 1, "id": "test1_blank", "type": "baseline", "checked": "verified", "max_length": 10}
|
||||
@@ -0,0 +1,117 @@
|
||||
{"pdf": "multi_column_miss.pdf", "page": 1, "id": "multi_column_miss_00", "type": "present", "text": "Corporate social responsibility and the tobacco industry: hope or hype?"}
|
||||
{"pdf": "multi_column_miss.pdf", "page": 1, "id": "multi_column_miss_01", "type": "present", "text": "this leaves BAT to argue why it should not be held to be largely accountable for the annual deaths of some 754 600 smokers, and Philip Morris some 803 600 smokers."}
|
||||
{"pdf": "multi_column_miss.pdf", "page": 1, "id": "multi_column_miss_02", "type": "present", "text": "The term \"corporate social responsibility\" is in vogue at the moment but as a concept it is vague and means different things to different people.", "max_diffs": 2}
|
||||
{"pdf": "multi_column_miss.pdf", "page": 1, "id": "multi_column_miss_03", "type": "present", "text": "Over the past three decades increasing pressure from non-governmental"}
|
||||
{"pdf": "multi_column_miss.pdf", "page": 1, "id": "multi_column_miss_04", "type": "absent", "text": "Downloaded from http://tobaccocontrol.bmj.com/"}
|
||||
|
||||
{"pdf": "multi_column_miss.pdf", "page": 1, "id": "multi_column_miss_10", "type": "order", "before": "Corporate social responsibility and the tobacco industry: hope or hype?", "after": "The unprecedented expansion of power and influence of TNCs over the past three decades has accelerated global trade and development, but also environmental damage and abuses of", "max_diffs": 2}
|
||||
{"pdf": "multi_column_miss.pdf", "page": 1, "id": "multi_column_miss_11", "type": "order", "before": "It now looks like that with vigilance", "after": "this leaves BAT to argue why it should not be held to be largely accountable for the annual deaths", "max_diffs": 2}
|
||||
{"pdf": "multi_column_miss.pdf", "page": 1, "id": "multi_column_miss_12", "type": "order", "before": "Corporate social responsibility (CSR) emerged from a realisation among transnational corporations", "after": " perspective on its own behaviour; and reflects on whether marketing tobacco is antithetical to social responsibility.", "max_diffs": 2}
|
||||
|
||||
{"pdf": "discoverworld_crazy_table4.pdf", "page": 1, "id": "discoverworld_crazy_table4_00", "type": "present", "text": "Table 4: Baseline model performance on each of the three scoring metrics"}
|
||||
{"pdf": "discoverworld_crazy_table4.pdf", "page": 1, "id": "discoverworld_crazy_table4_01", "type": "present", "text": "Table 5: Baseline model performance on each of the three scoring metrics"}
|
||||
{"pdf": "discoverworld_crazy_table4.pdf", "page": 1, "id": "discoverworld_crazy_table4_02", "type": "present", "text": "We use the GPT-4O model for all our agents due to its higher performance and lower cost compared to other models. For space we provide"}
|
||||
|
||||
{"pdf": "mattsnotes.pdf", "page": 1, "id": "mattsnotes_minediff_00", "type": "present", "checked": "verified", "text": "The-Stack-V2"}
|
||||
{"pdf": "mattsnotes.pdf", "page": 1, "id": "mattsnotes_minediff_01", "type": "present", "checked": "verified", "text": "SE, whatever we've scraped"}
|
||||
{"pdf": "mattsnotes.pdf", "page": 1, "id": "mattsnotes_minediff_02", "type": "present", "checked": "verified", "text": "HQ DCLM"}
|
||||
{"pdf": "mattsnotes.pdf", "page": 2, "id": "mattsnotes_minediff_03", "type": "present", "checked": "verified", "text": "Order by repo"}
|
||||
{"pdf": "mattsnotes.pdf", "page": 3, "id": "mattsnotes_minediff_04", "type": "present", "checked": "verified", "text": "ARCH + TRAINING"}
|
||||
|
||||
{"pdf": "buildingnotes.pdf", "page": 1, "id": "building_notes_00", "type": "present", "checked": "verified", "text": "Master Bath", "case_sensitive": false}
|
||||
{"pdf": "buildingnotes.pdf", "page": 1, "id": "building_notes_01", "type": "present", "checked": "verified", "text": "Laundry", "case_sensitive": false}
|
||||
{"pdf": "buildingnotes.pdf", "page": 1, "id": "building_notes_02", "type": "present", "checked": "verified", "text": "Guest Bath", "case_sensitive": false}
|
||||
|
||||
{"pdf": "lincoln_letter.pdf", "page": 1, "id": "lincoln_letter_minediff_00", "type": "present", "checked": "verified", "text": "January 10th 1864."}
|
||||
{"pdf": "lincoln_letter.pdf", "page": 1, "id": "lincoln_letter_minediff_01", "type": "present", "checked": "verified", "text": "Major General Hitchcock, Commissioner of Exchanges, is authorized and directed to offer Brigadier General Trimble, now a prisoner of war in Fort McHenry, in exchange for Major White, who is held as a prisoner at Richmond."}
|
||||
{"pdf": "lincoln_letter.pdf", "page": 1, "id": "lincoln_letter_minediff_03", "type": "present", "checked": "verified", "text": "He is also directed to send forward the offer of exchange by Henry M. Warfield, Esq. of Baltimore, under a flag of truce, and give him a pass to City Point."}
|
||||
|
||||
{"pdf": "openstax_caculus_pg_273.pdf", "page": 1, "id": "openstax_caculus_pg_273_minediff_02", "type": "present", "checked": "verified", "text": "Use the graph of the position function to determine the time intervals when the velocity is positive, negative, or zero."}
|
||||
{"pdf": "openstax_caculus_pg_273.pdf", "page": 1, "id": "openstax_caculus_pg_273_minediff_03", "type": "present", "checked": "verified", "text": "Use the graph of the velocity function to determine the time intervals when the acceleration is positive, negative, or zero."}
|
||||
{"pdf": "openstax_caculus_pg_273.pdf", "page": 1, "id": "openstax_caculus_pg_273_minediff_04", "type": "order", "before": "150.", "after": "157."}
|
||||
{"pdf": "openstax_caculus_pg_273.pdf", "page": 1, "id": "openstax_caculus_pg_273_minediff_05", "type": "order", "before": "150.", "after": "158."}
|
||||
{"pdf": "openstax_caculus_pg_273.pdf", "page": 1, "id": "openstax_caculus_pg_273_minediff_06", "type": "order", "before": "150.", "after": "159."}
|
||||
|
||||
{"pdf": "multi_column_miss.pdf", "page": 1, "id": "multi_column_miss_minediff_01", "type": "present", "checked": "verified", "text": "This report first provides the context and development of CSR; then, from internal company documents, examines how PM came to its own version."}
|
||||
{"pdf": "multi_column_miss.pdf", "page": 1, "id": "multi_column_miss_minediff_02", "type": "present", "checked": "verified", "text": "This paper examines whether a tobacco company espousing CSR should be judged simply as a corporate entity along standards of business ethics, or as an irretrievably negative force in the realm of public health, thereby rendering CSR an oxymoron."}
|
||||
|
||||
{"pdf": "olmo2-pg4.pdf", "page": 1, "id": "olmo2-pg4_minediff_00", "type": "present", "checked": "verified", "text": "Table 1 Composition of the pretraining data for OLMo 2."}
|
||||
|
||||
{"pdf": "olmo2-pg4.pdf", "page": 1, "id": "olmo2-pg4_table00", "type": "table", "cell": "Type"}
|
||||
{"pdf": "olmo2-pg4.pdf", "page": 1, "id": "olmo2-pg4_table01", "type": "table", "cell": "3.32T", "left": "3.71T"}
|
||||
{"pdf": "olmo2-pg4.pdf", "page": 1, "id": "olmo2-pg4_table02", "type": "table", "cell": "3.32T", "right": "21.32T"}
|
||||
{"pdf": "olmo2-pg4.pdf", "page": 1, "id": "olmo2-pg4_table03", "type": "table", "cell": "11.8B", "up": "12.2B"}
|
||||
{"pdf": "olmo2-pg4.pdf", "page": 1, "id": "olmo2-pg4_table04", "type": "table", "cell": "11.8B", "down": "3.7B"}
|
||||
{"pdf": "olmo2-pg4.pdf", "page": 1, "id": "olmo2-pg4_table05", "type": "table", "cell": "3.32T", "top_heading": "Words"}
|
||||
{"pdf": "olmo2-pg4.pdf", "page": 1, "id": "olmo2-pg4_table06", "type": "table", "cell": "arXiv", "top_heading": "Source"}
|
||||
{"pdf": "olmo2-pg4.pdf", "page": 1, "id": "olmo2-pg4_table07", "type": "table", "cell": "47.2B", "top_heading": "Bytes"}
|
||||
{"pdf": "olmo2-pg4.pdf", "page": 1, "id": "olmo2-pg4_table08", "type": "table", "cell": "Math proofs code", "left_heading": "Algebraic Stack"}
|
||||
|
||||
{"pdf": "discoverworld_crazy_table4.pdf", "page": 1, "id": "olmo2-discoverworld_crazy_table4_t00", "type": "table", "cell": "Quadratic regression", "left": "Challenge"}
|
||||
{"pdf": "discoverworld_crazy_table4.pdf", "page": 1, "id": "olmo2-discoverworld_crazy_table4_t01", "type": "table", "cell": "Instrument Use", "left": "Normal"}
|
||||
{"pdf": "discoverworld_crazy_table4.pdf", "page": 1, "id": "olmo2-discoverworld_crazy_table4_t02", "type": "table", "cell": "0.87", "top_heading": "Procedure"}
|
||||
{"pdf": "discoverworld_crazy_table4.pdf", "page": 1, "id": "olmo2-discoverworld_crazy_table4_t03", "type": "table", "cell": "0.87", "top_heading": "ReACT"}
|
||||
{"pdf": "discoverworld_crazy_table4.pdf", "page": 1, "id": "olmo2-discoverworld_crazy_table4_t04", "type": "table", "cell": "Pick-and-place object", "left_heading": "27"}
|
||||
{"pdf": "discoverworld_crazy_table4.pdf", "page": 1, "id": "olmo2-discoverworld_crazy_table4_t05", "type": "table", "cell": "0.66", "right": "0.44"}
|
||||
{"pdf": "discoverworld_crazy_table4.pdf", "page": 1, "id": "olmo2-discoverworld_crazy_table4_t06", "type": "table", "cell": "Interact with a moving agent", "top_heading": "Unit Test Topic"}
|
||||
|
||||
{"pdf": "earnings.pdf", "page": 1, "id": "earnings_table00", "type": "table", "cell": "1,136", "top_heading": "Year Ended"}
|
||||
{"pdf": "earnings.pdf", "page": 1, "id": "earnings_table01", "type": "table", "cell": "Year Ended"}
|
||||
{"pdf": "earnings.pdf", "page": 1, "id": "earnings_table02", "type": "table", "cell": "680", "up": "1,892"}
|
||||
{"pdf": "earnings.pdf", "page": 1, "id": "earnings_table03", "type": "table", "cell": "2,532", "left_heading": "Research and development"}
|
||||
{"pdf": "earnings.pdf", "page": 1, "id": "earnings_table04", "type": "absent", "text": "62"}
|
||||
|
||||
|
||||
{"pdf": "mathfuncs.pdf", "page": 1, "id": "mathfuncs_00", "type": "order", "before": "Euler's Identity", "after": "Pythagorean Theorem"}
|
||||
{"pdf": "mathfuncs.pdf", "page": 1, "id": "mathfuncs_01", "type": "order", "before": "Pythagorean Theorem", "after": "The Fundamental Theorem of Calculus"}
|
||||
{"pdf": "mathfuncs.pdf", "page": 1, "id": "mathfuncs_02", "type": "order", "before": "The Fundamental Theorem of Calculus", "after": "Maxwell's Equations"}
|
||||
{"pdf": "mathfuncs.pdf", "page": 1, "id": "mathfuncs_03", "type": "math", "math": "e^{i \\pi}+1=0"}
|
||||
{"pdf": "mathfuncs.pdf", "page": 1, "id": "mathfuncs_04", "type": "math", "math": "\\int_{a}^{b} f(x) d x=F(b)-F(a)"}
|
||||
{"pdf": "mathfuncs.pdf", "page": 1, "id": "mathfuncs_05", "type": "math", "math": "a^{2}+b^{2}=c^{2}"}
|
||||
{"pdf": "mathfuncs.pdf", "page": 1, "id": "mathfuncs_06", "type": "math", "math": "\\nabla \\times \\mathbf{E}=-\\frac{\\partial \\mathbf{B}}{\\partial t}"}
|
||||
|
||||
{"pdf": "mathfuncs_colswitch.pdf", "page": 1, "id": "mathfuncscol_00", "type": "order", "before": "Euler's Identity", "after": "Pythagorean Theorem"}
|
||||
{"pdf": "mathfuncs_colswitch.pdf", "page": 1, "id": "mathfuncscol_01", "type": "order", "before": "Pythagorean Theorem", "after": "The Fundamental Theorem of Calculus"}
|
||||
{"pdf": "mathfuncs_colswitch.pdf", "page": 1, "id": "mathfuncscol_02", "type": "order", "before": "The Fundamental Theorem of Calculus", "after": "Maxwell's Equations"}
|
||||
{"pdf": "mathfuncs_colswitch.pdf", "page": 1, "id": "mathfuncscol_03", "type": "math", "math": "e^{i \\pi}+1=0"}
|
||||
{"pdf": "mathfuncs_colswitch.pdf", "page": 1, "id": "mathfuncscol_04", "type": "math", "math": "\\int_{a}^{b} f(x) d x=F(b)-F(a)"}
|
||||
{"pdf": "mathfuncs_colswitch.pdf", "page": 1, "id": "mathfuncscol_05", "type": "math", "math": "a^{2}+b^{2}=c^{2}"}
|
||||
{"pdf": "mathfuncs_colswitch.pdf", "page": 1, "id": "mathfuncscol_06", "type": "math", "math": "\\nabla \\times \\mathbf{E}=-\\frac{\\partial \\mathbf{B}}{\\partial t}"}
|
||||
|
||||
{"pdf": "math_2503_04086.pdf", "page": 1, "id": "math_2503_04086_00", "type": "math", "math": "\\lambda_{g}=\\sum_{s \\in S} \\zeta_{n}^{\\psi(g s)}=\\sum_{i=1}^{k}\\left[\\sum_{s, R s=\\mathcal{I}_{i}} \\zeta_{n}^{\\psi(g s)}\\right]"}
|
||||
{"pdf": "math_2503_04086.pdf", "page": 1, "id": "math_2503_04086_01", "type": "math", "math": "\\lambda_{g}=\\lambda_{g^{\\prime}}"}
|
||||
{"pdf": "math_2503_04086.pdf", "page": 1, "id": "math_2503_04086_02", "type": "math", "math": "u \\in\\left(R / \\operatorname{Ann}_{R}\\left(x_{i}\\right)\\right)^{\\times}"}
|
||||
{"pdf": "math_2503_04086.pdf", "page": 1, "id": "math_2503_04086_03", "type": "math", "math": "\\lambda_{g}=\\sum_{i=1}^{k} c\\left(g, R / \\operatorname{Ann}_{R}\\left(x_{i}\\right)\\right)"}
|
||||
{"pdf": "math_2503_04086.pdf", "page": 1, "id": "math_2503_04086_04", "type": "present", "text": "We also thank Ján Mináč for his constant encouragement and support."}
|
||||
{"pdf": "math_2503_04086.pdf", "page": 1, "id": "math_2503_04086_05", "type": "present", "text": "Allgemeine theorie der Gaußschen Summen in endlichen kommutativen Ringe"}
|
||||
|
||||
{"pdf": "test-graphical-text.pdf", "page": 1, "id": "test_graphical_text_00", "type": "present", "text": "THE VISION"}
|
||||
{"pdf": "test-graphical-text.pdf", "page": 1, "id": "test_graphical_text_01", "type": "present", "text": "by FutureSkill"}
|
||||
{"pdf": "test-graphical-text.pdf", "page": 1, "id": "test_graphical_text_02", "type": "present", "text": "THE POWER OF STORYTELLING"}
|
||||
|
||||
{"pdf": "small_page_size.pdf", "page": 1, "id": "small_page_size_00", "type": "present", "text": "Since the use of bones has, however, become general, the turnip crop has been, in many instances, ten-fold, and in few less than four or five-fold its former bulk"}
|
||||
{"pdf": "small_page_size.pdf", "page": 1, "id": "small_page_size_01", "type": "order", "before": "Earthy and saline matter", "after": "Cartilage and jelly"}
|
||||
{"pdf": "small_page_size.pdf", "page": 1, "id": "small_page_size_02", "type": "present", "text": "The cartilage, indeed, when the bones have been buried in a dry situation, is very indestructible"}
|
||||
{"pdf": "small_page_size.pdf", "page": 1, "id": "small_page_size_03", "type": "absent", "text": "BRITISH HUSBANDRY"}
|
||||
{"pdf": "small_page_size.pdf", "page": 1, "id": "small_page_size_04", "type": "absent", "text": "400"}
|
||||
{"pdf": "small_page_size.pdf", "page": 1, "id": "small_page_size_05", "type": "absent", "text": "XIX"}
|
||||
|
||||
{"pdf": "headers_footers/ff0f0b22c55d8b90dd77d153f48e144fc9db_pg2.pdf", "page": 1, "id": "ff0f0b22c55d8b90dd77d153f48e144fc9db_02a", "type": "absent", "text": "PLOS Neglected Tropical Diseases | www.plosntds.org"}
|
||||
{"pdf": "headers_footers/ff0f0b22c55d8b90dd77d153f48e144fc9db_pg2.pdf", "page": 1, "id": "ff0f0b22c55d8b90dd77d153f48e144fc9db_02b", "type": "absent", "text": "March 2014 | Volume 8 | Issue 3 | e2748"}
|
||||
{"pdf": "headers_footers/ff0f0b22c55d8b90dd77d153f48e144fc9db_pg2.pdf", "page": 1, "id": "ff0f0b22c55d8b90dd77d153f48e144fc9db_02c", "type": "absent", "text": "1", "last_n": 20}
|
||||
{"pdf": "headers_footers/ff0f0b22c55d8b90dd77d153f48e144fc9db_pg2.pdf", "page": 1, "id": "ff0f0b22c55d8b90dd77d153f48e144fc9db_02d", "type": "absent", "text": "OPEN ACCESS Freely available online"}
|
||||
{"pdf": "headers_footers/ff0f0b22c55d8b90dd77d153f48e144fc9db_pg2.pdf", "page": 1, "id": "ff0f0b22c55d8b90dd77d153f48e144fc9db_02e", "type": "absent", "text": "PLOS NEGLECTED TROPICAL DISEASES"}
|
||||
{"pdf": "headers_footers/ff1fc6a205ad039139ce566851b6b260c929_pg1.pdf", "page": 1, "id": "ff1fc6a205ad039139ce566851b6b260c929_01a", "type": "absent", "text": "TELEDYNE ENERGY SYSTEMS, INC. A Teledyne Technologies Company"}
|
||||
{"pdf": "headers_footers/ff1fc6a205ad039139ce566851b6b260c929_pg1.pdf", "page": 1, "id": "ff1fc6a205ad039139ce566851b6b260c929_01b", "type": "absent", "text": "AEROJET ROCKETDYNE"}
|
||||
{"pdf": "headers_footers/ff1fc6a205ad039139ce566851b6b260c929_pg1.pdf", "page": 1, "id": "ff1fc6a205ad039139ce566851b6b260c929_01c", "type": "absent", "text": "1", "last_n": 20}
|
||||
{"pdf": "headers_footers/ff3d6e051903fe5ca9bc172ece14964c5632_pg1.pdf", "page": 1, "id": "ff3d6e051903fe5ca9bc172ece14964c5632_01a", "type": "absent", "text": "farbod4ever@gmail.com :\u0646\u0648\u064a\u0633\u0646\u062f\u0647 \u0631\u0627\u0628\u0637", "max_diffs": 4}
|
||||
{"pdf": "headers_footers/ff3d6e051903fe5ca9bc172ece14964c5632_pg1.pdf", "page": 1, "id": "ff3d6e051903fe5ca9bc172ece14964c5632_01b", "type": "absent", "text": "Downloaded from jipm.irandoc.ac.ir at 6:51 IRST on Monday November 11th 2019"}
|
||||
{"pdf": "headers_footers/ff4f7dad78081cff727d19ab51c181d4a661_pg1.pdf", "page": 1, "id": "ff4f7dad78081cff727d19ab51c181d4a661_01a", "type": "absent", "text": "Download date: 28 Dec 2018"}
|
||||
{"pdf": "headers_footers/ff518b1240a66978f22035528ccb029450b5_pg2.pdf", "page": 1, "id": "ff518b1240a66978f22035528ccb029450b5_02a", "type": "absent", "text": "Woodworth et al.: Brief Notices"}
|
||||
{"pdf": "headers_footers/ff518b1240a66978f22035528ccb029450b5_pg2.pdf", "page": 1, "id": "ff518b1240a66978f22035528ccb029450b5_02b", "type": "absent", "text": "Published by BYU ScholarsArchive, 1997"}
|
||||
{"pdf": "headers_footers/ff518b1240a66978f22035528ccb029450b5_pg2.pdf", "page": 1, "id": "ff518b1240a66978f22035528ccb029450b5_02c", "type": "absent", "text": "199", "last_n": 20}
|
||||
{"pdf": "headers_footers/ffaac214730d2b8c2ec842e3618ccb9c4259_pg1.pdf", "page": 1, "id": "ffaac214730d2b8c2ec842e3618ccb9c4259_01a", "type": "absent", "text": "TILBURG UNIVERSITY"}
|
||||
{"pdf": "headers_footers/ffaac214730d2b8c2ec842e3618ccb9c4259_pg1.pdf", "page": 1, "id": "ffaac214730d2b8c2ec842e3618ccb9c4259_01b", "type": "absent", "text": "Download date: 15. Sep. 2023"}
|
||||
{"pdf": "headers_footers/fff590bed29a2854ac1f874dad5752ede1aa_pg1.pdf", "page": 1, "id": "fff590bed29a2854ac1f874dad5752ede1aa_01a", "type": "absent", "text": "Revision: 2.4"}
|
||||
{"pdf": "headers_footers/fff590bed29a2854ac1f874dad5752ede1aa_pg1.pdf", "page": 1, "id": "fff590bed29a2854ac1f874dad5752ede1aa_01b", "type": "absent", "text": "P/N 119-036"}
|
||||
{"pdf": "headers_footers/fff590bed29a2854ac1f874dad5752ede1aa_pg1.pdf", "page": 1, "id": "fff590bed29a2854ac1f874dad5752ede1aa_01c", "type": "absent", "text": "10 June 2019"}
|
||||
@@ -0,0 +1,11 @@
|
||||
Master - 7 1/4 - 36"
|
||||
Master Bath - 7 1/4 - 30"
|
||||
Laundry - 4 3/4 - 36"
|
||||
Bath - 7 1/4 - 24"
|
||||
MUD - 7 - 36"
|
||||
UTIL - 8 1/4 - 36"
|
||||
DOWN BATH - 7 1/4 - 32"
|
||||
BUT KIT - 6 3/4 - 30
|
||||
PANTRY - 4 3/4 - 24
|
||||
6 WEST - 32 9/8 - 32
|
||||
6 WEST BATH 5" - 24"
|
||||
@@ -0,0 +1,51 @@
|
||||
Table 4: Baseline model performance on each of the three scoring metrics (task completion, task process, explanatory knowledge discovery) across all 24 DISCOVERY WORLD tasks. Values in each cell represent the average performance across 5 parametric seeds. Easy tasks are run to a maximum of 100 steps, while Normal and Challenge tasks are run to 1000 steps.
|
||||
|
||||
| # | Topic | Task | ReACT | Plan+Execute | Hypothesizer |
|
||||
|---|---|---|---|---|---|
|
||||
| 1 | Proteomics | Clustering | 0.87 | 0.20 | 0.20 | 0.80 | 0.00 | 0.00 | 0.90 | 0.40 | 1.00 |
|
||||
| 2 | Chemistry | Exploring Combinations and Hill Climbing | 0.88 | 0.40 | 0.40 | 0.68 | 0.20 | 0.00 | 0.93 | 0.40 | 0.40 |
|
||||
| 3 | Archaeology | Correlations | 0.88 | 0.40 | 0.60 | 0.55 | 0.20 | 0.00 | 0.93 | 0.40 | 0.60 |
|
||||
| 4 | Easy | Single substances | 0.87 | 1.00 | 1.00 | 0.70 | 0.60 | 0.40 | 0.90 | 0.00 | 0.40 |
|
||||
| 5 | Normal | Mix of 3 substances | 0.82 | 0.00 | 0.00 | 0.87 | 0.40 | 0.00 | 0.93 | 0.60 | 0.40 |
|
||||
| 6 | Challenge | Mix of 4 substances | 0.90 | 0.40 | 0.00 | 0.90 | 0.40 | 0.00 | 0.97 | 0.00 | 0.00 |
|
||||
| 7 | Easy | Simple instrument | 0.27 | 0.60 | 0.00 | 0.33 | 0.20 | 0.00 | 0.60 | 0.20 | 0.50 |
|
||||
| 8 | Normal | Instrument Use | 0.72 | 0.40 | 0.30 | 0.74 | 0.00 | 0.00 | 0.64 | 0.40 | 0.40 |
|
||||
| 9 | Challenge | Correlation | 0.46 | 0.20 | 0.00 | 0.46 | 0.00 | 0.05 | 0.55 | 0.20 | 0.05 |
|
||||
| 10 | Reactor Lab | Regression | 0.42 | 0.00 | 0.40 | 0.44 | 0.00 | 0.10 | 0.38 | 0.00 | 0.20 |
|
||||
| 11 | Normal | Linear regression | 0.44 | 0.00 | 0.20 | 0.49 | 0.00 | 0.00 | 0.51 | 0.00 | 0.00 |
|
||||
| 12 | Challenge | Quadratic regression | 0.43 | 0.00 | 0.20 | 0.39 | 0.00 | 0.00 | 0.39 | 0.00 | 0.00 |
|
||||
| 13 | Easy | Simplified rules | 0.80 | 0.20 | 0.20 | 0.70 | 0.20 | 0.20 | 0.60 | 0.00 | 0.00 |
|
||||
| 14 | Normal | Presence rules | 0.91 | 0.60 | 0.00 | 0.84 | 0.40 | 0.00 | 0.56 | 0.00 | 0.00 |
|
||||
| 15 | Challenge | Logical Rules | 0.89 | 0.40 | 0.00 | 0.73 | 0.40 | 0.00 | 0.62 | 0.00 | 0.00 |
|
||||
| 16 | Easy | Open-ended discovery | 0.78 | 0.60 | 0.00 | 0.68 | 0.40 | 0.10 | 0.80 | 1.00 | 0.60 |
|
||||
| 17 | Normal | Multiple instruments | 0.58 | 0.00 | 0.13 | 0.45 | 0.00 | 0.13 | 0.16 | 0.00 | 0.33 |
|
||||
| 18 | Challenge | Novel instruments | 0.55 | 0.00 | 0.00 | 0.26 | 0.00 | 0.00 | 0.20 | 0.00 | 0.00 |
|
||||
| 19 | Easy | Look-up variables | 0.33 | 0.00 | 0.00 | 0.53 | 0.00 | 0.07 | 0.13 | 0.40 | 0.00 |
|
||||
| 20 | Normal | Measure 2 variables | 0.51 | 0.00 | 0.05 | 0.34 | 0.00 | 0.00 | 0.11 | 0.00 | 0.00 |
|
||||
| 21 | Challenge | Measure 5 variables | 0.43 | 0.00 | 0.00 | 0.15 | 0.00 | 0.00 | 0.22 | 0.00 | 0.03 |
|
||||
| 22 | Easy | Rosetta-stone style linguistic discovery of alien language | 0.40 | 0.40 | 0.20 | 0.30 | 0.00 | 0.00 | 0.20 | 0.20 | 0.00 |
|
||||
| 23 | Normal | Noun and verb | 0.20 | 0.00 | 0.00 | 0.68 | 0.40 | 0.00 | 0.84 | 0.40 | 0.00 |
|
||||
| 24 | Challenge | Noun, adj., and verb | 0.49 | 0.00 | 0.00 | 0.55 | 0.20 | 0.05 | 0.15 | 0.00 | 0.00 |
|
||||
| Average (Easy) | 0.59 | 0.38 | 0.25 | 0.56 | 0.18 | 0.11 | 0.56 | 0.28 | 0.34 |
|
||||
| Average (Normal) | 0.63 | 0.18 | 0.14 | 0.64 | 0.18 | 0.02 | 0.58 | 0.23 | 0.19 |
|
||||
| Average (Challenge) | 0.63 | 0.18 | 0.10 | 0.50 | 0.15 | 0.01 | 0.49 | 0.08 | 0.08 |
|
||||
|
||||
Table 5: Baseline model performance on each of the three scoring metrics (task completion, task process, explanatory knowledge discovery) across all 10 unit test tasks. Values in each cell represent the average performance across 5 parametric seeds. Unit tests tasks are run to a maximum of 100 steps.
|
||||
|
||||
| # | Unit Test Topic | ReACT | Plan+Execute | Hypothesizer |
|
||||
|---|---|---|---|---|
|
||||
| 25 | Multi-turn dialog with an agent | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 |
|
||||
| 26 | Measure an object with an instrument | 0.87 | 0.60 | 0.73 | 0.40 | 1.00 | 1.00 |
|
||||
| 27 | Pick-and-place object | 0.90 | 0.80 | 0.80 | 0.60 | 1.00 | 1.00 |
|
||||
| 28 | Pick-and-give object | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 | 1.00 |
|
||||
| 29 | Read DiscoveryFeed posts | 1.00 | 1.00 | 0.90 | 0.80 | 1.00 | 1.00 |
|
||||
| 30 | Move through doors | 0.58 | 0.20 | 0.25 | 0.00 | 0.30 | 0.00 |
|
||||
| 31 | Using keys with doors | 0.69 | 0.20 | 0.54 | 0.00 | 0.69 | 0.00 |
|
||||
| 32 | Navigate to a specific room in a house | 0.20 | 0.20 | 0.20 | 0.00 | 0.20 | 0.20 |
|
||||
| 33 | Search an environment for an object | 0.80 | 0.80 | 0.60 | 0.60 | 1.00 | 1.00 |
|
||||
| 34 | Interact with a moving agent | 0.60 | 0.20 | 0.53 | 0.00 | 0.53 | 0.20 |
|
||||
| Average (Unit Tests) | 0.76 | 0.60 | 0.66 | 0.44 | 0.77 | 0.64 |
|
||||
|
||||
4.2 Baseline Agent Models
|
||||
|
||||
The baseline agents are described below, with model performance on Discovery tasks shown in Table 4, and performance on Unit Tests shown in Table 5. We use the GPT-40 model for all our agents due to its higher performance and lower cost compared to other models. For space we provide
|
||||
@@ -0,0 +1,33 @@
|
||||
Recently Issued Accounting Pronouncements
|
||||
|
||||
Recently Adopted Accounting Pronouncement
|
||||
|
||||
In November 2023, the Financial Accounting Standards Board, or FASB, issued a new accounting standard requiring disclosures of significant expenses in operating segments. We adopted this standard in our fiscal year 2025 annual report. Refer to Note 16 of the Notes to the Consolidated Financial Statements in Part IV, Item 15 of this Annual Report on Form 10-K for further information.
|
||||
|
||||
Recent Accounting Pronouncements Not Yet Adopted
|
||||
|
||||
In December 2023, the FASB issued a new accounting standard which includes new and updated income tax disclosures, including disaggregation of information in the rate reconciliation and income taxes paid. We expect to adopt this standard in our fiscal year 2026 annual report. We do not expect the adoption of this standard to have a material impact on our Consolidated Financial Statements other than additional disclosures.
|
||||
|
||||
In November 2024, the FASB issued a new accounting standard requiring disclosures of certain additional expense information on an annual and interim basis, including, among other items, the amounts of purchases of inventory, employee compensation, depreciation and intangible asset amortization included within each income statement expense caption, as applicable. We expect to adopt this standard in our fiscal year 2028 annual report. We do not expect the adoption of this standard to have a material impact on our Consolidated Financial Statements other than additional disclosures.
|
||||
|
||||
Note 2 - Business Combination
|
||||
|
||||
Termination of the Arm Share Purchase Agreement
|
||||
|
||||
In February 2022, NVIDIA and SoftBank Group Corp, or SoftBank, announced the termination of the Share Purchase Agreement whereby NVIDIA would have acquired Arm from SoftBank. The parties agreed to terminate it due to significant regulatory challenges preventing the completion of the transaction. We recorded an acquisition termination cost of $1.4 billion in fiscal year 2023 reflecting the write-off of the prepayment provided at signing.
|
||||
|
||||
Note 3 - Stock-Based Compensation
|
||||
|
||||
Stock-based compensation expense is associated with RSUs, PSUs, market-based PSUs, and our ESPP.
|
||||
|
||||
Consolidated Statements of Income include stock-based compensation expense, net of amounts capitalized into inventory and subsequently recognized to cost of revenue, as follows:
|
||||
|
||||
| Year Ended | Jan 29, 2023 | Jan 28, 2024 | Jan 29, 2023 |
|
||||
|------------|--------------|--------------|--------------|
|
||||
| (In millions) | $138 | $141 | $138 |
|
||||
| Cost of revenue | $178 | $141 | $138 |
|
||||
| Research and development | 3,423 | 2,532 | 1,892 |
|
||||
| Sales, general and administrative | 1,136 | 876 | 680 |
|
||||
| Total | $4,737 | $3,549 | $2,710 |
|
||||
|
||||
Stock-based compensation capitalized in inventories was not significant during fiscal years 2025, 2024, and 2023.
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
Lassa Fever in Post-Conflict Sierra Leone
|
||||
|
||||
Jeffrey G. Shaffer1, Donald S. Grant2, John S. Schieffelin3, Matt L. Boisen4,5, Augustine Goba2,9, Jessica N. Hartnett4, Danielle C. Levy4, Rachael E. Yenni4, Lina M. Moses4,6, Mohammed Fullah2, Mambo Momoh2, Mbalu Fonnie2, Richard Fonnie2, Lansana Kanneh2, Veronica J. Koroma2, Kande Kargbo2, Darin Ottomassathien9, Ivana J. Muncy5, Abigail B. Jones5, Megan M. Illic7, Peter C. Kulakosky8, Allyson M. Haislip4, Christopher M. Bishop4, Deborah H. Elliot3, Bethany L. Brown5, Hu Zhu9, Kathryn M. Hastie10, Kristian G. Andersen11,12, Stephen K. Gire11,12, Shervin Tabrizi11,12, Ridhi Tariyal12, Mathew Stremlau11,12, Alex Matschiner7, Darryl B. Sampey7, Jennifer S. Spence6, Robert W. Cross4,13, Joan B. Geisbert13, Onikepe A. Folarin14, Christian T. Happi14, Kelly R. Pitts5, F. Jon Geske5, Thomas W. Geisbert13, Erica Ollmann Saphire9,15, James E. Robinson3, Russell B. Wilson8, Pardis C. Sabeti10,11,16, Lee A. Henderson8, S. Humarr Khan25, Daniel G. Bausch61, Luis M. Branco8,17, Robert F. Garry4,12,17*, the Viral Hemorrhagic Fever Consortium†
|
||||
|
||||
1 Department of Biostatistics and Bioinformatics, Tulane School of Public Health and Tropical Medicine, New Orleans, Louisiana, United States of America, 2 Lassa Fever Program, Kenema Government Hospital, Kenema, Sierra Leone, 3 Sections of Infectious Disease, Departments of Pediatrics and Internal Medicine, School of Medicine, Tulane University, New Orleans, Louisiana, United States of America, 4 Department of Microbiology and Immunology, Tulane University, New Orleans, Louisiana, United States of America, 5 Corgenix, Inc., Broomfield, Colorado, United States of America, 6 Department of Tropical Medicine, Tulane School of Public Health and Tropical Medicine, New Orleans, Louisiana, United States of America, 7 Biofactura Inc., Rockville, Maryland, United States of America, 8 Autoimmune Technologies, LLC, New Orleans, Louisiana, United States of America, 9 Vybion Inc., Ithaca, New York, United States of America, 10 Department of Immunology and Microbial Science, The Scripps Research Institute, La Jolla, California, United States of America, 11 FAS Center for Systems Biology, Department of Organismic and Evolutionary Biology, Harvard University, Cambridge, Massachusetts, United States of America, 12 Broad Institute, Cambridge, Massachusetts, United States of America, 13 Department of Microbiology and Immunology, University of Texas Medical Branch at Galveston, Galveston, Texas, United States of America, 14 Department of Biological Sciences, College of Natural Sciences, Redeemers University, Redemption City, Ogun State, Nigeria, 15 The Skaggs Institute for Chemical Biology, The Scripps Research Institute, La Jolla, California, United States of America, 16 Department of Immunology and Infectious Disease, Harvard School of Public Health, Boston, Massachusetts, United States of America, 17 Zalgen Labs, LLC, Germantown, Maryland, United States of America
|
||||
|
||||
Abstract
|
||||
|
||||
Background: Lassa fever (LF), an often-fatal hemorrhagic disease caused by Lassa virus (LASV), is a major public health threat in West Africa. When the violent civil conflict in Sierra Leone (1991 to 2002) ended, an international consortium assisted in restoration of the LF program at Kenema Government Hospital (KGH) in an area with the world’s highest incidence of the disease.
|
||||
|
||||
Methodology/Principal Findings: Clinical and laboratory records of patients presenting to the KGH Lassa Ward in the post-conflict period were organized electronically. Recombinant antigen-based LF immunoassays were used to assess LASV antigenemia and LASV-specific antibodies in patients who met criteria for suspected LF. KGH has been reestablished as a center for LF treatment and research, with over 500 suspected cases now presenting yearly. Higher case fatality rates (CFRs) in LF patients were observed compared to studies conducted prior to the civil conflict. Different criteria for defining LF stages and differences in sensitivity of assays likely account for these differences. The highest incidence of LF in Sierra Leone was observed during the dry season. LF cases were observed in ten of Sierra Leone’s thirteen districts, with numerous cases from outside the traditional endemic zone. Deaths in patients presenting with LASV antigenemia were skewed towards individuals less than 29 years of age. Women self-reporting as pregnant were significantly overrepresented among LASV antigenemic patients. The CFR of ribavirin-treated patients presenting early in acute infection was lower than in untreated subjects.
|
||||
|
||||
Conclusions/Significance: Lassa fever remains a major public health threat in Sierra Leone. Outreach activities should expand because LF may be more widespread in Sierra Leone than previously recognized. Enhanced case finding to ensure rapid diagnosis and treatment is imperative to reduce mortality. Even with ribavirin treatment, there was a high rate of fatalities underscoring the need to develop more effective and/or supplemental treatments for LF.
|
||||
|
||||
Citation: Shaffer JG, Grant DS, Schieffelin JS, Boisen ML, Goba A, et al. (2014) Lassa Fever in Post-Conflict Sierra Leone. PLoS Negl Trop Dis 8(3): e2748. doi:10.1371/journal.pntd.0002748
|
||||
|
||||
Editor: Brian Bird, Centers for Disease Control and Prevention, United States of America
|
||||
|
||||
Received August 10, 2013; Accepted February 3, 2014; Published March 20, 2014
|
||||
|
||||
Copyright: © 2014 Shaffer et al. This is an open-access article distributed under the terms of the Creative Commons Attribution License, which permits unrestricted use, distribution, and reproduction in any medium, provided the original author and source are credited.
|
||||
|
||||
Funding: This work was supported by National Institute of Allergy and Infectious Diseases grants/contracts AI067188, A082119, AI2008031, AI2009061, AI104216, AI104621, P20GM103501, HHSN272200900018C, HHSN272200900049C, HHSN272201000022C, and an Investigators in the Pathogenesis of Infectious Disease award from the Burroughs Wellcome Fund. The funders had no role in study design, data collection and analysis, decision to publish, or preparation of the manuscript.
|
||||
+10
@@ -0,0 +1,10 @@
|
||||
RTG Degradation Primer and Application to MMRTG
|
||||
|
||||
Nuclear and Emerging Technology for Space (NETS) 2015
|
||||
February 23-26, 2015
|
||||
Abstract 5107
|
||||
|
||||
Presenting Author: Tom Hammel, Teledyne Energy Systems
|
||||
Co-Authors: Russell Bennett, Teledyne Energy Systems
|
||||
Robert Sievers, Teledyne Energy Systems
|
||||
Bill Otting, Aerojet Rocketdyne
|
||||
+9
@@ -0,0 +1,9 @@
|
||||
بررسی دیدگاه و نظرات کتابداران و اعضای هیئت علمی دانشگاه شیراز در بره گیری از فناوری شبکههای پی سیم در کتابخانههای دانشگاهی
|
||||
|
||||
چکیده: نظر به اهمیت و کاربرد گسترده شبکههای پی سیم در محیطهای دانشگاهی در کشورهای پیشرفته و استفاده از آن در خدمات کتابخانهای، بهرهگیری از این فناوری در کتابخانههای دانشگاهی کشور احساس میشود. این پژوهش با استفاده از روش پیامبرداری، با هدف معرفی شبکههای پی سیم، به بررسی دیدگاه و نظرات کتابداران و اعضای هیئت علمی دانشگاه شیراز، در یک کارگیری از این شبکهها در کتابخانههای دانشگاهی پرداخت. باقیماندههای تحقیق نشان داد که کتابداران تا مایل زیادی به استفاده از شبکههای پی سیم در امر خدمات کتابخانههای تظیم میباشند و قفسهخوانی و دسترسی به فهرست عمومی پوسته دارند. و گسترش گی پوشش و سیاستهای پیشرفته و سیاستهای پیشرفته در کتابخانهها را خواستار شدهاند. در حالی که اعضای هیئت علمی، ضرورت استفاده بیشتر از این شبکهها در کل محیط دانشگاه و دسترسی به منابع کتابخانهای از خارج از محیط کتابخانه را خواستار شدهاند. در کل، با توجه به نتایج حاصل از این تحقیق میتوان گفت که هر چند استفاده از شبکههای پی سیم در کتابخانهها پیشبینیهای زیادی و نوآوریهای میباشد و هنوز در کشور ما چندان مورد توجه قرار نگرفته و ناشناخته مانده است، اما تماشای کتابداران، پژوهشگران و استادان به استفاده و کاربرد آن در محیطهای دانشگاهی زیاد است.
|
||||
|
||||
کلیدواژهها: شبکههای پی سیم؛ کتابخانههای دانشگاهی؛ اعضای هیئت علمی؛ کتابداران؛ رایانههای قابل حمل؛ رایانههای دستی؛ منابع و خدمات کتابخانهای؛ دانشگاه شیراز
|
||||
|
||||
farbod4ever@gmail.com
|
||||
|
||||
نویسنده رابطه:
|
||||
+11
@@ -0,0 +1,11 @@
|
||||
Molecular markers of breast cancer metastasis
|
||||
Weigelt, B.
|
||||
|
||||
Citation for published version (APA):
|
||||
Weigelt, B. (2005). Molecular markers of breast cancer metastasis
|
||||
|
||||
General rights
|
||||
It is not permitted to download or to forward/distribute the text or part of it without the consent of the author(s) and/or copyright holder(s), other than for strictly personal, individual use, unless the work is under an open content license (like Creative Commons).
|
||||
|
||||
Disclaimer/Complaints regulations
|
||||
If you believe that digital publication of certain material infringes any of your rights or (privacy) interests, please let the Library know, stating your reasons. In case of a legitimate complaint, the Library will make the material inaccessible and/or remove it from the website. Please Ask the Library: http://uba.uva.nl/en/contact, or a letter to: Library of the University of Amsterdam, Secretariat, Singel 425, 1012 WP Amsterdam, The Netherlands. You will be contacted as soon as possible.
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user