Add all files and directories

.gitignore

@@ -0,0 +1,168 @@
+private.py
+.DS_Store
+local.env
+experiments
+test_data
+training
+wandb
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/

CLA.md

@@ -0,0 +1,24 @@
+Marker Contributor Agreement
+
+This Marker Contributor Agreement ("MCA") applies to any contribution that you make to any product or project managed by us (the "project"), and sets out the intellectual property rights you grant to us in the contributed materials. The term "us" shall mean Vikas Paruchuri. The term "you" shall mean the person or entity identified below. 
+
+If you agree to be bound by these terms, sign by writing "I have read the CLA document and I hereby sign the CLA" in response to the CLA bot Github comment. Read this agreement carefully before signing. These terms and conditions constitute a binding legal agreement.
+
+1. The term 'contribution' or 'contributed materials' means any source code, object code, patch, tool, sample, graphic, specification, manual, documentation, or any other material posted or submitted by you to the project. 
+2. With respect to any worldwide copyrights, or copyright applications and registrations, in your contribution: 
+   - you hereby assign to us joint ownership, and to the extent that such assignment is or becomes invalid, ineffective or unenforceable, you hereby grant to us a perpetual, irrevocable, non-exclusive, worldwide, no-charge, royalty free, unrestricted license to exercise all rights under those copyrights. This includes, at our option, the right to sublicense these same rights to third parties through multiple levels of sublicensees or other licensing arrangements, including dual-license structures for commercial customers; 
+   - you agree that each of us can do all things in relation to your contribution as if each of us were the sole owners, and if one of us makes a derivative work of your contribution, the one who makes the derivative work (or has it made will be the sole owner of that derivative work; 
+   - you agree that you will not assert any moral rights in your contribution against us, our licensees or transferees; 
+   - you agree that we may register a copyright in your contribution and exercise all ownership rights associated with it; and 
+   - you agree that neither of us has any duty to consult with, obtain the consent of, pay or render an accounting to the other for any use or distribution of vour contribution. 
+3. With respect to any patents you own, or that you can license without payment to any third party, you hereby grant to us a perpetual, irrevocable, non-exclusive, worldwide, no-charge, royalty-free license to:
+   - make, have made, use, sell, offer to sell, import, and otherwise transfer your contribution in whole or in part, alone or in combination with or included in any product, work or materials arising out of the project to which your contribution was submitted, and
+   - at our option, to sublicense these same rights to third parties through multiple levels of sublicensees or other licensing arrangements. 
+If you or your affiliates institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the contribution or any project it was submitted to constitutes direct or contributory patent infringement, then any patent licenses granted to you under this agreement for that contribution shall terminate as of the date such litigation is filed.
+4. Except as set out above, you keep all right, title, and interest in your contribution. The rights that you grant to us under these terms are effective on the date you first submitted a contribution to us, even if your submission took place before the date you sign these terms. Any contribution we make available under any license will also be made available under a suitable FSF (Free Software Foundation) or OSI (Open Source Initiative) approved license. 
+5. You covenant, represent, warrant and agree that: 
+   - each contribution that you submit is and shall be an original work of authorship and you can legally grant the rights set out in this MCA; 
+   - to the best of your knowledge, each contribution will not violate any third party's copyrights, trademarks, patents, or other intellectual property rights; and 
+   - each contribution shall be in compliance with U.S. export control laws and other applicable export and import laws.
+You agree to notify us if you become aware of any circumstance which would make any of the foregoing representations inaccurate in any respect. Vikas Paruchuri may publicly disclose your participation in the project, including the fact that you have signed the MCA. 
+6. This MCA is governed by the laws of the State of California and applicable U.S. Federal law. Any choice of law rules will not apply.

LICENSE

@@ -0,0 +1,674 @@
+                    GNU GENERAL PUBLIC LICENSE
+                       Version 3, 29 June 2007
+
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
+
+                            Preamble
+
+  The GNU General Public License is a free, copyleft license for
+software and other kinds of works.
+
+  The licenses for most software and other practical works are designed
+to take away your freedom to share and change the works.  By contrast,
+the GNU General Public License is intended to guarantee your freedom to
+share and change all versions of a program--to make sure it remains free
+software for all its users.  We, the Free Software Foundation, use the
+GNU General Public License for most of our software; it applies also to
+any other work released this way by its authors.  You can apply it to
+your programs, too.
+
+  When we speak of free software, we are referring to freedom, not
+price.  Our General Public Licenses are designed to make sure that you
+have the freedom to distribute copies of free software (and charge for
+them if you wish), that you receive source code or can get it if you
+want it, that you can change the software or use pieces of it in new
+free programs, and that you know you can do these things.
+
+  To protect your rights, we need to prevent others from denying you
+these rights or asking you to surrender the rights.  Therefore, you have
+certain responsibilities if you distribute copies of the software, or if
+you modify it: responsibilities to respect the freedom of others.
+
+  For example, if you distribute copies of such a program, whether
+gratis or for a fee, you must pass on to the recipients the same
+freedoms that you received.  You must make sure that they, too, receive
+or can get the source code.  And you must show them these terms so they
+know their rights.
+
+  Developers that use the GNU GPL protect your rights with two steps:
+(1) assert copyright on the software, and (2) offer you this License
+giving you legal permission to copy, distribute and/or modify it.
+
+  For the developers' and authors' protection, the GPL clearly explains
+that there is no warranty for this free software.  For both users' and
+authors' sake, the GPL requires that modified versions be marked as
+changed, so that their problems will not be attributed erroneously to
+authors of previous versions.
+
+  Some devices are designed to deny users access to install or run
+modified versions of the software inside them, although the manufacturer
+can do so.  This is fundamentally incompatible with the aim of
+protecting users' freedom to change the software.  The systematic
+pattern of such abuse occurs in the area of products for individuals to
+use, which is precisely where it is most unacceptable.  Therefore, we
+have designed this version of the GPL to prohibit the practice for those
+products.  If such problems arise substantially in other domains, we
+stand ready to extend this provision to those domains in future versions
+of the GPL, as needed to protect the freedom of users.
+
+  Finally, every program is threatened constantly by software patents.
+States should not allow patents to restrict development and use of
+software on general-purpose computers, but in those that do, we wish to
+avoid the special danger that patents applied to a free program could
+make it effectively proprietary.  To prevent this, the GPL assures that
+patents cannot be used to render the program non-free.
+
+  The precise terms and conditions for copying, distribution and
+modification follow.
+
+                       TERMS AND CONDITIONS
+
+  0. Definitions.
+
+  "This License" refers to version 3 of the GNU General Public License.
+
+  "Copyright" also means copyright-like laws that apply to other kinds of
+works, such as semiconductor masks.
+
+  "The Program" refers to any copyrightable work licensed under this
+License.  Each licensee is addressed as "you".  "Licensees" and
+"recipients" may be individuals or organizations.
+
+  To "modify" a work means to copy from or adapt all or part of the work
+in a fashion requiring copyright permission, other than the making of an
+exact copy.  The resulting work is called a "modified version" of the
+earlier work or a work "based on" the earlier work.
+
+  A "covered work" means either the unmodified Program or a work based
+on the Program.
+
+  To "propagate" a work means to do anything with it that, without
+permission, would make you directly or secondarily liable for
+infringement under applicable copyright law, except executing it on a
+computer or modifying a private copy.  Propagation includes copying,
+distribution (with or without modification), making available to the
+public, and in some countries other activities as well.
+
+  To "convey" a work means any kind of propagation that enables other
+parties to make or receive copies.  Mere interaction with a user through
+a computer network, with no transfer of a copy, is not conveying.
+
+  An interactive user interface displays "Appropriate Legal Notices"
+to the extent that it includes a convenient and prominently visible
+feature that (1) displays an appropriate copyright notice, and (2)
+tells the user that there is no warranty for the work (except to the
+extent that warranties are provided), that licensees may convey the
+work under this License, and how to view a copy of this License.  If
+the interface presents a list of user commands or options, such as a
+menu, a prominent item in the list meets this criterion.
+
+  1. Source Code.
+
+  The "source code" for a work means the preferred form of the work
+for making modifications to it.  "Object code" means any non-source
+form of a work.
+
+  A "Standard Interface" means an interface that either is an official
+standard defined by a recognized standards body, or, in the case of
+interfaces specified for a particular programming language, one that
+is widely used among developers working in that language.
+
+  The "System Libraries" of an executable work include anything, other
+than the work as a whole, that (a) is included in the normal form of
+packaging a Major Component, but which is not part of that Major
+Component, and (b) serves only to enable use of the work with that
+Major Component, or to implement a Standard Interface for which an
+implementation is available to the public in source code form.  A
+"Major Component", in this context, means a major essential component
+(kernel, window system, and so on) of the specific operating system
+(if any) on which the executable work runs, or a compiler used to
+produce the work, or an object code interpreter used to run it.
+
+  The "Corresponding Source" for a work in object code form means all
+the source code needed to generate, install, and (for an executable
+work) run the object code and to modify the work, including scripts to
+control those activities.  However, it does not include the work's
+System Libraries, or general-purpose tools or generally available free
+programs which are used unmodified in performing those activities but
+which are not part of the work.  For example, Corresponding Source
+includes interface definition files associated with source files for
+the work, and the source code for shared libraries and dynamically
+linked subprograms that the work is specifically designed to require,
+such as by intimate data communication or control flow between those
+subprograms and other parts of the work.
+
+  The Corresponding Source need not include anything that users
+can regenerate automatically from other parts of the Corresponding
+Source.
+
+  The Corresponding Source for a work in source code form is that
+same work.
+
+  2. Basic Permissions.
+
+  All rights granted under this License are granted for the term of
+copyright on the Program, and are irrevocable provided the stated
+conditions are met.  This License explicitly affirms your unlimited
+permission to run the unmodified Program.  The output from running a
+covered work is covered by this License only if the output, given its
+content, constitutes a covered work.  This License acknowledges your
+rights of fair use or other equivalent, as provided by copyright law.
+
+  You may make, run and propagate covered works that you do not
+convey, without conditions so long as your license otherwise remains
+in force.  You may convey covered works to others for the sole purpose
+of having them make modifications exclusively for you, or provide you
+with facilities for running those works, provided that you comply with
+the terms of this License in conveying all material for which you do
+not control copyright.  Those thus making or running the covered works
+for you must do so exclusively on your behalf, under your direction
+and control, on terms that prohibit them from making any copies of
+your copyrighted material outside their relationship with you.
+
+  Conveying under any other circumstances is permitted solely under
+the conditions stated below.  Sublicensing is not allowed; section 10
+makes it unnecessary.
+
+  3. Protecting Users' Legal Rights From Anti-Circumvention Law.
+
+  No covered work shall be deemed part of an effective technological
+measure under any applicable law fulfilling obligations under article
+11 of the WIPO copyright treaty adopted on 20 December 1996, or
+similar laws prohibiting or restricting circumvention of such
+measures.
+
+  When you convey a covered work, you waive any legal power to forbid
+circumvention of technological measures to the extent such circumvention
+is effected by exercising rights under this License with respect to
+the covered work, and you disclaim any intention to limit operation or
+modification of the work as a means of enforcing, against the work's
+users, your or third parties' legal rights to forbid circumvention of
+technological measures.
+
+  4. Conveying Verbatim Copies.
+
+  You may convey verbatim copies of the Program's source code as you
+receive it, in any medium, provided that you conspicuously and
+appropriately publish on each copy an appropriate copyright notice;
+keep intact all notices stating that this License and any
+non-permissive terms added in accord with section 7 apply to the code;
+keep intact all notices of the absence of any warranty; and give all
+recipients a copy of this License along with the Program.
+
+  You may charge any price or no price for each copy that you convey,
+and you may offer support or warranty protection for a fee.
+
+  5. Conveying Modified Source Versions.
+
+  You may convey a work based on the Program, or the modifications to
+produce it from the Program, in the form of source code under the
+terms of section 4, provided that you also meet all of these conditions:
+
+    a) The work must carry prominent notices stating that you modified
+    it, and giving a relevant date.
+
+    b) The work must carry prominent notices stating that it is
+    released under this License and any conditions added under section
+    7.  This requirement modifies the requirement in section 4 to
+    "keep intact all notices".
+
+    c) You must license the entire work, as a whole, under this
+    License to anyone who comes into possession of a copy.  This
+    License will therefore apply, along with any applicable section 7
+    additional terms, to the whole of the work, and all its parts,
+    regardless of how they are packaged.  This License gives no
+    permission to license the work in any other way, but it does not
+    invalidate such permission if you have separately received it.
+
+    d) If the work has interactive user interfaces, each must display
+    Appropriate Legal Notices; however, if the Program has interactive
+    interfaces that do not display Appropriate Legal Notices, your
+    work need not make them do so.
+
+  A compilation of a covered work with other separate and independent
+works, which are not by their nature extensions of the covered work,
+and which are not combined with it such as to form a larger program,
+in or on a volume of a storage or distribution medium, is called an
+"aggregate" if the compilation and its resulting copyright are not
+used to limit the access or legal rights of the compilation's users
+beyond what the individual works permit.  Inclusion of a covered work
+in an aggregate does not cause this License to apply to the other
+parts of the aggregate.
+
+  6. Conveying Non-Source Forms.
+
+  You may convey a covered work in object code form under the terms
+of sections 4 and 5, provided that you also convey the
+machine-readable Corresponding Source under the terms of this License,
+in one of these ways:
+
+    a) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by the
+    Corresponding Source fixed on a durable physical medium
+    customarily used for software interchange.
+
+    b) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by a
+    written offer, valid for at least three years and valid for as
+    long as you offer spare parts or customer support for that product
+    model, to give anyone who possesses the object code either (1) a
+    copy of the Corresponding Source for all the software in the
+    product that is covered by this License, on a durable physical
+    medium customarily used for software interchange, for a price no
+    more than your reasonable cost of physically performing this
+    conveying of source, or (2) access to copy the
+    Corresponding Source from a network server at no charge.
+
+    c) Convey individual copies of the object code with a copy of the
+    written offer to provide the Corresponding Source.  This
+    alternative is allowed only occasionally and noncommercially, and
+    only if you received the object code with such an offer, in accord
+    with subsection 6b.
+
+    d) Convey the object code by offering access from a designated
+    place (gratis or for a charge), and offer equivalent access to the
+    Corresponding Source in the same way through the same place at no
+    further charge.  You need not require recipients to copy the
+    Corresponding Source along with the object code.  If the place to
+    copy the object code is a network server, the Corresponding Source
+    may be on a different server (operated by you or a third party)
+    that supports equivalent copying facilities, provided you maintain
+    clear directions next to the object code saying where to find the
+    Corresponding Source.  Regardless of what server hosts the
+    Corresponding Source, you remain obligated to ensure that it is
+    available for as long as needed to satisfy these requirements.
+
+    e) Convey the object code using peer-to-peer transmission, provided
+    you inform other peers where the object code and Corresponding
+    Source of the work are being offered to the general public at no
+    charge under subsection 6d.
+
+  A separable portion of the object code, whose source code is excluded
+from the Corresponding Source as a System Library, need not be
+included in conveying the object code work.
+
+  A "User Product" is either (1) a "consumer product", which means any
+tangible personal property which is normally used for personal, family,
+or household purposes, or (2) anything designed or sold for incorporation
+into a dwelling.  In determining whether a product is a consumer product,
+doubtful cases shall be resolved in favor of coverage.  For a particular
+product received by a particular user, "normally used" refers to a
+typical or common use of that class of product, regardless of the status
+of the particular user or of the way in which the particular user
+actually uses, or expects or is expected to use, the product.  A product
+is a consumer product regardless of whether the product has substantial
+commercial, industrial or non-consumer uses, unless such uses represent
+the only significant mode of use of the product.
+
+  "Installation Information" for a User Product means any methods,
+procedures, authorization keys, or other information required to install
+and execute modified versions of a covered work in that User Product from
+a modified version of its Corresponding Source.  The information must
+suffice to ensure that the continued functioning of the modified object
+code is in no case prevented or interfered with solely because
+modification has been made.
+
+  If you convey an object code work under this section in, or with, or
+specifically for use in, a User Product, and the conveying occurs as
+part of a transaction in which the right of possession and use of the
+User Product is transferred to the recipient in perpetuity or for a
+fixed term (regardless of how the transaction is characterized), the
+Corresponding Source conveyed under this section must be accompanied
+by the Installation Information.  But this requirement does not apply
+if neither you nor any third party retains the ability to install
+modified object code on the User Product (for example, the work has
+been installed in ROM).
+
+  The requirement to provide Installation Information does not include a
+requirement to continue to provide support service, warranty, or updates
+for a work that has been modified or installed by the recipient, or for
+the User Product in which it has been modified or installed.  Access to a
+network may be denied when the modification itself materially and
+adversely affects the operation of the network or violates the rules and
+protocols for communication across the network.
+
+  Corresponding Source conveyed, and Installation Information provided,
+in accord with this section must be in a format that is publicly
+documented (and with an implementation available to the public in
+source code form), and must require no special password or key for
+unpacking, reading or copying.
+
+  7. Additional Terms.
+
+  "Additional permissions" are terms that supplement the terms of this
+License by making exceptions from one or more of its conditions.
+Additional permissions that are applicable to the entire Program shall
+be treated as though they were included in this License, to the extent
+that they are valid under applicable law.  If additional permissions
+apply only to part of the Program, that part may be used separately
+under those permissions, but the entire Program remains governed by
+this License without regard to the additional permissions.
+
+  When you convey a copy of a covered work, you may at your option
+remove any additional permissions from that copy, or from any part of
+it.  (Additional permissions may be written to require their own
+removal in certain cases when you modify the work.)  You may place
+additional permissions on material, added by you to a covered work,
+for which you have or can give appropriate copyright permission.
+
+  Notwithstanding any other provision of this License, for material you
+add to a covered work, you may (if authorized by the copyright holders of
+that material) supplement the terms of this License with terms:
+
+    a) Disclaiming warranty or limiting liability differently from the
+    terms of sections 15 and 16 of this License; or
+
+    b) Requiring preservation of specified reasonable legal notices or
+    author attributions in that material or in the Appropriate Legal
+    Notices displayed by works containing it; or
+
+    c) Prohibiting misrepresentation of the origin of that material, or
+    requiring that modified versions of such material be marked in
+    reasonable ways as different from the original version; or
+
+    d) Limiting the use for publicity purposes of names of licensors or
+    authors of the material; or
+
+    e) Declining to grant rights under trademark law for use of some
+    trade names, trademarks, or service marks; or
+
+    f) Requiring indemnification of licensors and authors of that
+    material by anyone who conveys the material (or modified versions of
+    it) with contractual assumptions of liability to the recipient, for
+    any liability that these contractual assumptions directly impose on
+    those licensors and authors.
+
+  All other non-permissive additional terms are considered "further
+restrictions" within the meaning of section 10.  If the Program as you
+received it, or any part of it, contains a notice stating that it is
+governed by this License along with a term that is a further
+restriction, you may remove that term.  If a license document contains
+a further restriction but permits relicensing or conveying under this
+License, you may add to a covered work material governed by the terms
+of that license document, provided that the further restriction does
+not survive such relicensing or conveying.
+
+  If you add terms to a covered work in accord with this section, you
+must place, in the relevant source files, a statement of the
+additional terms that apply to those files, or a notice indicating
+where to find the applicable terms.
+
+  Additional terms, permissive or non-permissive, may be stated in the
+form of a separately written license, or stated as exceptions;
+the above requirements apply either way.
+
+  8. Termination.
+
+  You may not propagate or modify a covered work except as expressly
+provided under this License.  Any attempt otherwise to propagate or
+modify it is void, and will automatically terminate your rights under
+this License (including any patent licenses granted under the third
+paragraph of section 11).
+
+  However, if you cease all violation of this License, then your
+license from a particular copyright holder is reinstated (a)
+provisionally, unless and until the copyright holder explicitly and
+finally terminates your license, and (b) permanently, if the copyright
+holder fails to notify you of the violation by some reasonable means
+prior to 60 days after the cessation.
+
+  Moreover, your license from a particular copyright holder is
+reinstated permanently if the copyright holder notifies you of the
+violation by some reasonable means, this is the first time you have
+received notice of violation of this License (for any work) from that
+copyright holder, and you cure the violation prior to 30 days after
+your receipt of the notice.
+
+  Termination of your rights under this section does not terminate the
+licenses of parties who have received copies or rights from you under
+this License.  If your rights have been terminated and not permanently
+reinstated, you do not qualify to receive new licenses for the same
+material under section 10.
+
+  9. Acceptance Not Required for Having Copies.
+
+  You are not required to accept this License in order to receive or
+run a copy of the Program.  Ancillary propagation of a covered work
+occurring solely as a consequence of using peer-to-peer transmission
+to receive a copy likewise does not require acceptance.  However,
+nothing other than this License grants you permission to propagate or
+modify any covered work.  These actions infringe copyright if you do
+not accept this License.  Therefore, by modifying or propagating a
+covered work, you indicate your acceptance of this License to do so.
+
+  10. Automatic Licensing of Downstream Recipients.
+
+  Each time you convey a covered work, the recipient automatically
+receives a license from the original licensors, to run, modify and
+propagate that work, subject to this License.  You are not responsible
+for enforcing compliance by third parties with this License.
+
+  An "entity transaction" is a transaction transferring control of an
+organization, or substantially all assets of one, or subdividing an
+organization, or merging organizations.  If propagation of a covered
+work results from an entity transaction, each party to that
+transaction who receives a copy of the work also receives whatever
+licenses to the work the party's predecessor in interest had or could
+give under the previous paragraph, plus a right to possession of the
+Corresponding Source of the work from the predecessor in interest, if
+the predecessor has it or can get it with reasonable efforts.
+
+  You may not impose any further restrictions on the exercise of the
+rights granted or affirmed under this License.  For example, you may
+not impose a license fee, royalty, or other charge for exercise of
+rights granted under this License, and you may not initiate litigation
+(including a cross-claim or counterclaim in a lawsuit) alleging that
+any patent claim is infringed by making, using, selling, offering for
+sale, or importing the Program or any portion of it.
+
+  11. Patents.
+
+  A "contributor" is a copyright holder who authorizes use under this
+License of the Program or a work on which the Program is based.  The
+work thus licensed is called the contributor's "contributor version".
+
+  A contributor's "essential patent claims" are all patent claims
+owned or controlled by the contributor, whether already acquired or
+hereafter acquired, that would be infringed by some manner, permitted
+by this License, of making, using, or selling its contributor version,
+but do not include claims that would be infringed only as a
+consequence of further modification of the contributor version.  For
+purposes of this definition, "control" includes the right to grant
+patent sublicenses in a manner consistent with the requirements of
+this License.
+
+  Each contributor grants you a non-exclusive, worldwide, royalty-free
+patent license under the contributor's essential patent claims, to
+make, use, sell, offer for sale, import and otherwise run, modify and
+propagate the contents of its contributor version.
+
+  In the following three paragraphs, a "patent license" is any express
+agreement or commitment, however denominated, not to enforce a patent
+(such as an express permission to practice a patent or covenant not to
+sue for patent infringement).  To "grant" such a patent license to a
+party means to make such an agreement or commitment not to enforce a
+patent against the party.
+
+  If you convey a covered work, knowingly relying on a patent license,
+and the Corresponding Source of the work is not available for anyone
+to copy, free of charge and under the terms of this License, through a
+publicly available network server or other readily accessible means,
+then you must either (1) cause the Corresponding Source to be so
+available, or (2) arrange to deprive yourself of the benefit of the
+patent license for this particular work, or (3) arrange, in a manner
+consistent with the requirements of this License, to extend the patent
+license to downstream recipients.  "Knowingly relying" means you have
+actual knowledge that, but for the patent license, your conveying the
+covered work in a country, or your recipient's use of the covered work
+in a country, would infringe one or more identifiable patents in that
+country that you have reason to believe are valid.
+
+  If, pursuant to or in connection with a single transaction or
+arrangement, you convey, or propagate by procuring conveyance of, a
+covered work, and grant a patent license to some of the parties
+receiving the covered work authorizing them to use, propagate, modify
+or convey a specific copy of the covered work, then the patent license
+you grant is automatically extended to all recipients of the covered
+work and works based on it.
+
+  A patent license is "discriminatory" if it does not include within
+the scope of its coverage, prohibits the exercise of, or is
+conditioned on the non-exercise of one or more of the rights that are
+specifically granted under this License.  You may not convey a covered
+work if you are a party to an arrangement with a third party that is
+in the business of distributing software, under which you make payment
+to the third party based on the extent of your activity of conveying
+the work, and under which the third party grants, to any of the
+parties who would receive the covered work from you, a discriminatory
+patent license (a) in connection with copies of the covered work
+conveyed by you (or copies made from those copies), or (b) primarily
+for and in connection with specific products or compilations that
+contain the covered work, unless you entered into that arrangement,
+or that patent license was granted, prior to 28 March 2007.
+
+  Nothing in this License shall be construed as excluding or limiting
+any implied license or other defenses to infringement that may
+otherwise be available to you under applicable patent law.
+
+  12. No Surrender of Others' Freedom.
+
+  If conditions are imposed on you (whether by court order, agreement or
+otherwise) that contradict the conditions of this License, they do not
+excuse you from the conditions of this License.  If you cannot convey a
+covered work so as to satisfy simultaneously your obligations under this
+License and any other pertinent obligations, then as a consequence you may
+not convey it at all.  For example, if you agree to terms that obligate you
+to collect a royalty for further conveying from those to whom you convey
+the Program, the only way you could satisfy both those terms and this
+License would be to refrain entirely from conveying the Program.
+
+  13. Use with the GNU Affero General Public License.
+
+  Notwithstanding any other provision of this License, you have
+permission to link or combine any covered work with a work licensed
+under version 3 of the GNU Affero General Public License into a single
+combined work, and to convey the resulting work.  The terms of this
+License will continue to apply to the part which is the covered work,
+but the special requirements of the GNU Affero General Public License,
+section 13, concerning interaction through a network will apply to the
+combination as such.
+
+  14. Revised Versions of this License.
+
+  The Free Software Foundation may publish revised and/or new versions of
+the GNU General Public License from time to time.  Such new versions will
+be similar in spirit to the present version, but may differ in detail to
+address new problems or concerns.
+
+  Each version is given a distinguishing version number.  If the
+Program specifies that a certain numbered version of the GNU General
+Public License "or any later version" applies to it, you have the
+option of following the terms and conditions either of that numbered
+version or of any later version published by the Free Software
+Foundation.  If the Program does not specify a version number of the
+GNU General Public License, you may choose any version ever published
+by the Free Software Foundation.
+
+  If the Program specifies that a proxy can decide which future
+versions of the GNU General Public License can be used, that proxy's
+public statement of acceptance of a version permanently authorizes you
+to choose that version for the Program.
+
+  Later license versions may give you additional or different
+permissions.  However, no additional obligations are imposed on any
+author or copyright holder as a result of your choosing to follow a
+later version.
+
+  15. Disclaimer of Warranty.
+
+  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
+APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
+HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
+OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
+THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
+IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
+ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
+
+  16. Limitation of Liability.
+
+  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
+WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
+THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
+GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
+USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
+DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
+PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
+EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGES.
+
+  17. Interpretation of Sections 15 and 16.
+
+  If the disclaimer of warranty and limitation of liability provided
+above cannot be given local legal effect according to their terms,
+reviewing courts shall apply local law that most closely approximates
+an absolute waiver of all civil liability in connection with the
+Program, unless a warranty or assumption of liability accompanies a
+copy of the Program in return for a fee.
+
+                     END OF TERMS AND CONDITIONS
+
+            How to Apply These Terms to Your New Programs
+
+  If you develop a new program, and you want it to be of the greatest
+possible use to the public, the best way to achieve this is to make it
+free software which everyone can redistribute and change under these terms.
+
+  To do so, attach the following notices to the program.  It is safest
+to attach them to the start of each source file to most effectively
+state the exclusion of warranty; and each file should have at least
+the "copyright" line and a pointer to where the full notice is found.
+
+    Marker pdf to markdown converter
+    Copyright (C) 2023  Vikas Paruchuri
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+Also add information on how to contact you by electronic and paper mail.
+
+  If the program does terminal interaction, make it output a short
+notice like this when it starts in an interactive mode:
+
+    Marker  Copyright (C) 2023  Vikas Paruchuri
+    This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
+    This is free software, and you are welcome to redistribute it
+    under certain conditions; type `show c' for details.
+
+The hypothetical commands `show w' and `show c' should show the appropriate
+parts of the General Public License.  Of course, your program's commands
+might be different; for a GUI interface, you would use an "about box".
+
+  You should also get your employer (if you work as a programmer) or school,
+if any, to sign a "copyright disclaimer" for the program, if necessary.
+For more information on this, and how to apply and follow the GNU GPL, see
+<https://www.gnu.org/licenses/>.
+
+  The GNU General Public License does not permit incorporating your program
+into proprietary programs.  If your program is a subroutine library, you
+may consider it more useful to permit linking proprietary applications with
+the library.  If this is what you want to do, use the GNU Lesser General
+Public License instead of this License.  But first, please read
+<https://www.gnu.org/licenses/why-not-lgpl.html>.

benchmark.py

@@ -0,0 +1,157 @@
+import argparse
+import tempfile
+import time
+from collections import defaultdict
+
+from tqdm import tqdm
+import pypdfium2 as pdfium
+
+from marker.convert import convert_single_pdf
+from marker.logger import configure_logging
+from marker.models import load_all_models
+from marker.benchmark.scoring import score_text
+from marker.pdf.extract_text import naive_get_text
+import json
+import os
+import subprocess
+import shutil
+from tabulate import tabulate
+import torch
+
+configure_logging()
+
+
+def start_memory_profiling():
+    torch.cuda.memory._record_memory_history(
+        max_entries=100000
+    )
+
+
+def stop_memory_profiling(memory_file):
+    try:
+        torch.cuda.memory._dump_snapshot(memory_file)
+    except Exception as e:
+        logger.error(f"Failed to capture memory snapshot {e}")
+
+        # Stop recording memory snapshot history.
+    torch.cuda.memory._record_memory_history(enabled=None)
+
+
+def nougat_prediction(pdf_filename, batch_size=1):
+    out_dir = tempfile.mkdtemp()
+    subprocess.run(["nougat", pdf_filename, "-o", out_dir, "--no-skipping", "--recompute", "--batchsize", str(batch_size)], check=True)
+    md_file = os.listdir(out_dir)[0]
+    with open(os.path.join(out_dir, md_file), "r") as f:
+        data = f.read()
+    shutil.rmtree(out_dir)
+    return data
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Benchmark PDF to MD conversion.  Needs source pdfs, and a refernece folder with the correct markdown.")
+    parser.add_argument("in_folder", help="Input PDF files")
+    parser.add_argument("reference_folder", help="Reference folder with reference markdown files")
+    parser.add_argument("out_file", help="Output filename")
+    parser.add_argument("--nougat", action="store_true", help="Run nougat and compare", default=False)
+    # Nougat batch size 1 uses about as much VRAM as default marker settings
+    parser.add_argument("--marker_batch_multiplier", type=int, default=1, help="Batch size multiplier to use for marker when making predictions.")
+    parser.add_argument("--nougat_batch_size", type=int, default=1, help="Batch size to use for nougat when making predictions.")
+    parser.add_argument("--md_out_path", type=str, default=None, help="Output path for generated markdown files")
+    parser.add_argument("--profile_memory", action="store_true", help="Profile memory usage", default=False)
+
+    args = parser.parse_args()
+
+    methods = ["marker"]
+    if args.nougat:
+        methods.append("nougat")
+
+    if args.profile_memory:
+        start_memory_profiling()
+
+    model_lst = load_all_models()
+
+    if args.profile_memory:
+        stop_memory_profiling("model_load.pickle")
+
+    scores = defaultdict(dict)
+    benchmark_files = os.listdir(args.in_folder)
+    benchmark_files = [b for b in benchmark_files if b.endswith(".pdf")]
+    times = defaultdict(dict)
+    pages = defaultdict(int)
+
+    for idx, fname in tqdm(enumerate(benchmark_files)):
+        md_filename = fname.rsplit(".", 1)[0] + ".md"
+
+        reference_filename = os.path.join(args.reference_folder, md_filename)
+        with open(reference_filename, "r", encoding="utf-8") as f:
+            reference = f.read()
+
+        pdf_filename = os.path.join(args.in_folder, fname)
+        doc = pdfium.PdfDocument(pdf_filename)
+        pages[fname] = len(doc)
+
+        for method in methods:
+            start = time.time()
+            if method == "marker":
+                if args.profile_memory:
+                    start_memory_profiling()
+                full_text, _, out_meta = convert_single_pdf(pdf_filename, model_lst, batch_multiplier=args.marker_batch_multiplier)
+                if args.profile_memory:
+                    stop_memory_profiling(f"marker_memory_{idx}.pickle")
+            elif method == "nougat":
+                full_text = nougat_prediction(pdf_filename, batch_size=args.nougat_batch_size)
+            elif method == "naive":
+                full_text = naive_get_text(doc)
+            else:
+                raise ValueError(f"Unknown method {method}")
+
+            times[method][fname] = time.time() - start
+
+            score = score_text(full_text, reference)
+            scores[method][fname] = score
+
+            if args.md_out_path:
+                md_out_filename = f"{method}_{md_filename}"
+                with open(os.path.join(args.md_out_path, md_out_filename), "w+") as f:
+                    f.write(full_text)
+
+    total_pages = sum(pages.values())
+    with open(args.out_file, "w+") as f:
+        write_data = defaultdict(dict)
+        for method in methods:
+            total_time = sum(times[method].values())
+            file_stats = {
+                fname:
+                {
+                    "time": times[method][fname],
+                    "score": scores[method][fname],
+                    "pages": pages[fname]
+                }
+
+                for fname in benchmark_files
+            }
+            write_data[method] = {
+                "files": file_stats,
+                "avg_score": sum(scores[method].values()) / len(scores[method]),
+                "time_per_page": total_time / total_pages,
+                "time_per_doc": total_time / len(scores[method])
+            }
+
+        json.dump(write_data, f, indent=4)
+
+    summary_table = []
+    score_table = []
+    score_headers = benchmark_files
+    for method in methods:
+        summary_table.append([method, write_data[method]["avg_score"], write_data[method]["time_per_page"], write_data[method]["time_per_doc"]])
+        score_table.append([method, *[write_data[method]["files"][h]["score"] for h in score_headers]])
+
+    print(tabulate(summary_table, headers=["Method", "Average Score", "Time per page", "Time per document"]))
+    print("")
+    print("Scores by file")
+    print(tabulate(score_table, headers=["Method", *score_headers]))
+
+
+if __name__ == "__main__":
+    main()
+

chunk_convert.py

@@ -0,0 +1,22 @@
+import argparse
+import subprocess
+import pkg_resources
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert a folder of PDFs to a folder of markdown files in chunks.")
+    parser.add_argument("in_folder", help="Input folder with pdfs.")
+    parser.add_argument("out_folder", help="Output folder")
+    args = parser.parse_args()
+
+    script_path = pkg_resources.resource_filename(__name__, 'chunk_convert.sh')
+
+    # Construct the command
+    cmd = f"{script_path} {args.in_folder} {args.out_folder}"
+
+    # Execute the shell script
+    subprocess.run(cmd, shell=True, check=True)
+
+
+if __name__ == "__main__":
+    main()

chunk_convert.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+
+trap 'pkill -P $$' SIGINT
+
+# Check if NUM_DEVICES is set
+if [[ -z "$NUM_DEVICES" ]]; then
+    echo "Please set the NUM_DEVICES environment variable."
+    exit 1
+fi
+
+if [[ -z "$NUM_WORKERS" ]]; then
+    echo "Please set the NUM_WORKERS environment variable."
+    exit 1
+fi
+
+
+# Get input folder and output folder from args
+if [[ -z "$1" ]]; then
+    echo "Please provide an input folder."
+    exit 1
+fi
+
+if [[ -z "$2" ]]; then
+    echo "Please provide an output folder."
+    exit 1
+fi
+
+INPUT_FOLDER=$1
+OUTPUT_FOLDER=$2
+
+# Loop from 0 to NUM_DEVICES and run the Python script in parallel
+for (( i=0; i<$NUM_DEVICES; i++ )); do
+    DEVICE_NUM=$i
+    export DEVICE_NUM
+    export NUM_DEVICES
+    export NUM_WORKERS
+    echo "Running convert.py on GPU $DEVICE_NUM"
+    cmd="CUDA_VISIBLE_DEVICES=$DEVICE_NUM marker $INPUT_FOLDER $OUTPUT_FOLDER --num_chunks $NUM_DEVICES --chunk_idx $DEVICE_NUM --workers $NUM_WORKERS"
+    [[ -n "$METADATA_FILE" ]] && cmd="$cmd --metadata_file $METADATA_FILE"
+    [[ -n "$MIN_LENGTH" ]] && cmd="$cmd --min_length $MIN_LENGTH"
+    eval $cmd &
+
+    sleep 5
+done
+
+# Wait for all background processes to finish
+wait

convert.py

@@ -0,0 +1,131 @@
+import argparse
+import os
+from typing import Dict, Optional
+
+import ray
+from tqdm import tqdm
+import math
+
+from marker.convert import convert_single_pdf
+from marker.output import markdown_exists, save_markdown
+from marker.pdf.utils import find_filetype
+from marker.pdf.extract_text import get_length_of_text
+from marker.models import load_all_models
+from marker.settings import settings
+from marker.logger import configure_logging
+import traceback
+import json
+
+configure_logging()
+
+
+@ray.remote(num_cpus=settings.RAY_CORES_PER_WORKER, num_gpus=.05 if settings.CUDA else 0)
+def process_single_pdf(filepath: str, out_folder: str, model_refs, metadata: Optional[Dict] = None, min_length: Optional[int] = None):
+    fname = os.path.basename(filepath)
+    if markdown_exists(out_folder, fname):
+        return
+
+    try:
+        # Skip trying to convert files that don't have a lot of embedded text
+        # This can indicate that they were scanned, and not OCRed properly
+        # Usually these files are not recent/high-quality
+        if min_length:
+            filetype = find_filetype(filepath)
+            if filetype == "other":
+                return 0
+
+            length = get_length_of_text(filepath)
+            if length < min_length:
+                return
+
+        full_text, images, out_metadata = convert_single_pdf(filepath, model_refs, metadata=metadata)
+        if len(full_text.strip()) > 0:
+            save_markdown(out_folder, fname, full_text, images, out_metadata)
+        else:
+            print(f"Empty file: {filepath}.  Could not convert.")
+    except Exception as e:
+        print(f"Error converting {filepath}: {e}")
+        print(traceback.format_exc())
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert multiple pdfs to markdown.")
+    parser.add_argument("in_folder", help="Input folder with pdfs.")
+    parser.add_argument("out_folder", help="Output folder")
+    parser.add_argument("--chunk_idx", type=int, default=0, help="Chunk index to convert")
+    parser.add_argument("--num_chunks", type=int, default=1, help="Number of chunks being processed in parallel")
+    parser.add_argument("--max", type=int, default=None, help="Maximum number of pdfs to convert")
+    parser.add_argument("--workers", type=int, default=5, help="Number of worker processes to use")
+    parser.add_argument("--metadata_file", type=str, default=None, help="Metadata json file to use for filtering")
+    parser.add_argument("--min_length", type=int, default=None, help="Minimum length of pdf to convert")
+
+    args = parser.parse_args()
+
+    in_folder = os.path.abspath(args.in_folder)
+    out_folder = os.path.abspath(args.out_folder)
+    files = [os.path.join(in_folder, f) for f in os.listdir(in_folder)]
+    files = [f for f in files if os.path.isfile(f)]
+    os.makedirs(out_folder, exist_ok=True)
+
+    # Handle chunks if we're processing in parallel
+    # Ensure we get all files into a chunk
+    chunk_size = math.ceil(len(files) / args.num_chunks)
+    start_idx = args.chunk_idx * chunk_size
+    end_idx = start_idx + chunk_size
+    files_to_convert = files[start_idx:end_idx]
+
+    # Limit files converted if needed
+    if args.max:
+        files_to_convert = files_to_convert[:args.max]
+
+    metadata = {}
+    if args.metadata_file:
+        metadata_file = os.path.abspath(args.metadata_file)
+        with open(metadata_file, "r") as f:
+            metadata = json.load(f)
+
+    total_processes = min(len(files_to_convert), args.workers)
+
+    ray.init(
+        num_cpus=total_processes,
+        num_gpus=1 if settings.CUDA else 0,
+        storage=settings.RAY_CACHE_PATH,
+        _temp_dir=settings.RAY_CACHE_PATH,
+        log_to_driver=settings.DEBUG
+    )
+
+    model_lst = load_all_models()
+    model_refs = ray.put(model_lst)
+
+    # Dynamically set GPU allocation per task based on GPU ram
+    gpu_frac = settings.VRAM_PER_TASK / settings.INFERENCE_RAM if settings.CUDA else 0
+
+    print(f"Converting {len(files_to_convert)} pdfs in chunk {args.chunk_idx + 1}/{args.num_chunks} with {total_processes} processes, and storing in {out_folder}")
+    futures = [
+        process_single_pdf.options(num_gpus=gpu_frac).remote(
+            filepath,
+            out_folder,
+            model_refs,
+            metadata=metadata.get(os.path.basename(filepath)),
+            min_length=args.min_length
+        ) for filepath in files_to_convert
+    ]
+
+    # Run all ray conversion tasks
+    progress_bar = tqdm(total=len(futures))
+    while len(futures) > 0:
+        finished, futures = ray.wait(
+            futures, timeout=7.0
+        )
+        finished_lst = ray.get(finished)
+        if isinstance(finished_lst, list):
+            progress_bar.update(len(finished_lst))
+        else:
+            progress_bar.update(1)
+
+    # Shutdown ray to free resources
+    ray.shutdown()
+
+
+if __name__ == "__main__":
+    main()

convert_single.py

@@ -0,0 +1,35 @@
+import argparse
+import os
+
+from marker.convert import convert_single_pdf
+from marker.logger import configure_logging
+from marker.models import load_all_models
+
+from marker.output import save_markdown
+
+configure_logging()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("filename", help="PDF file to parse")
+    parser.add_argument("output", help="Output base folder path")
+    parser.add_argument("--max_pages", type=int, default=None, help="Maximum number of pages to parse")
+    parser.add_argument("--langs", type=str, help="Languages to use for OCR, comma separated", default=None)
+    parser.add_argument("--batch_multiplier", type=int, default=2, help="How much to increase batch sizes")
+    args = parser.parse_args()
+
+    langs = args.langs.split(",") if args.langs else None
+
+    fname = args.filename
+    model_lst = load_all_models()
+    full_text, images, out_meta = convert_single_pdf(fname, model_lst, max_pages=args.max_pages, langs=langs, batch_multiplier=args.batch_multiplier)
+
+    fname = os.path.basename(fname)
+    subfolder_path = save_markdown(args.output, fname, full_text, images, out_meta)
+
+    print(f"Saved markdown to the {subfolder_path} folder")
+
+
+if __name__ == "__main__":
+    main()

data/.gitignore

@@ -0,0 +1,3 @@
+latex
+pdfs
+references

data/examples/marker/multicolcnn.md

@@ -0,0 +1,350 @@
+# An Aggregated Multicolumn Dilated Convolution Network For Perspective-Free Counting
+
+Diptodip Deb Georgia Institute of Technology [email protected] Jonathan Ventura University of Colorado Colorado Springs [email protected]
+
+## Abstract
+
+We propose the use of dilated filters to construct an aggregation module in a multicolumn convolutional neural network for perspective-free counting. Counting is a common problem in computer vision (e.g. traffic on the street or pedestrians in a crowd). Modern approaches to the counting problem involve the production of a density map via regression whose integral is equal to the number of objects in the image. However, objects in the image can occur at different scales (e.g. due to perspective effects) which can make it difficult for a learning agent to learn the proper density map. While the use of multiple columns to extract multiscale information from images has been shown before, our approach aggregates the multiscale information gathered by the multicolumn convolutional neural network to improve performance. Our experiments show that our proposed network outperforms the state-of-the-art on many benchmark datasets, and also that using our aggregation module in combination with a higher number of columns is beneficial for multiscale counting.
+
+## 1. Introduction
+
+Learning to count the number of objects in an image is a deceptively difficult problem with many interesting applications, such as surveillance [20], traffic monitoring [14] and medical image analysis [22]. In many of these application areas, the objects to be counted vary widely in appearance, size and shape, and labeled training data is typically sparse.
+
+These factors pose a significant computer vision and machine learning challenge.
+
+Lempitsky et al. [15] showed that it is possible to learn to count without learning to explicitly detect and localize individual objects. Instead, they propose learning to predict a density map whose integral over the image equals the number of objects in the image. This approach has been adopted by many later works (Cf. [18, 28]).
+
+However, in many counting problems, such as those counting cells in a microscope image, pedestrians in a crowd, or vehicles in a traffic jam, regressors trained on a single image scale are not reliable [18]. This is due to a variety of challenges including overlap of objects and perspective effects which cause significant variance in object shape, size and appearance.
+
+The most successful recent approaches address this issue by explicitly incorporating multi-scale information in the network [18,28]. These approaches either combine multiple networks which take input patches of different sizes [18]
+or combine multiple filtering paths ("columns") which have different size filters [28].
+
+Following on the intuition that multiscale integration is key to achieving good counting performance, we propose to incorporate dilated filters [25] into a multicolumn convolutional neural network design [28]. Dilated filters exponentially increase the network's receptive field without an exponential increase in parameters, allowing for efficient use of multiscale information. Convolutional neural networks with dilated filters have proven to provide competitive performance in image segmentation where multiscale analysis is also critical [25, 26]. By incorporating dilated filters into the multicolumn network design, we greatly increase the ability of the network to selectively aggregate multiscale information, without the need for explicit perspective maps during training and testing. We propose the
+"aggregated multicolumn dilated convolution network" or AMDCN which uses dilations to aggregate multiscale information. Our extensive experimental evaluation shows that this proposed network outperforms previous methods on many benchmark datasets.
+
+## 2. Related Work
+
+Counting using a supervised regressor to formulate a density map was first shown by [15]. In this paper, Lempitsky et al. show that the minimal annotation of a single dot blurred by a Gaussian kernel produces a sufficient density map to train a network to count. All of the counting methods that we examine as well as the method we use in
+
+![1_image_0.png](1_image_0.png)
+
+D
+our paper follow this method of producing a density map via regression. This is particularly advantageous because a sufficiently accurate regressor can also locate the objects in the image via this method. However, the Lempitsky paper ignores the issue of perspective scaling and other scaling issues. The work of [27] introduces CNNs (convolutional neural networks) for the purposes of crowd counting, but performs regression on similarly scaled image patches.
+
+These issues are addressed by the work of [18]. Rubio et al. show that a fully convolutional neural network can be used to produce a supervised regressor that produces density maps as in [15]. They further demonstrate a method dubbed HydraCNN which essentially combines multiple convolutional networks that take in differently scaled image patches in order to incorporate multiscale, global information from the image. The premise of this method is that a single regressor will fail to accurately represent the difference in values of the features of an image caused by perspective shifts (scaling effects) [18].
+
+However, the architectures of both [18] and [27] are not fully convolutional due to requiring multiple image patches and, as discussed in [25], the experiments of [11, 17] and [9, 12, 16] leave it unclear as to whether rescaling patches of the image is truly necessary in order to solve dense prediction problems via convolutional neural networks. Moreover, these approaches seem to saturate in performance at three columns, which means the network is extracting information from fewer scales. The work of [25] proposes the use of dilated convolutions as a simpler alternative that does not require sampling of rescaled image patches to provide global, scale-aware information to the network. A fully convolutional approach to multiscale counting has been proposed by [28], in which a multicolumn convolutional network gathers features of different scales by using convolutions of increasing kernel sizes from column to column instead of scaling image patches. Further, DeepLab has used dilated convolutions in multiple columns to extract scale information for segmentation [8]. We build on these approaches with our aggregator module as described in Section 3.1, which should allow for extracting information from more scales.
+
+It should be noted that other methods of counting exist, including training a network to recognize deep object features via only providing the counts of the objects of interest in an image [21] and using CNNs (convolutional neural networks) along with boosting in order to improve the results
+
+![2_image_0.png](2_image_0.png)
+
+of regression for production of density maps [24]. In the same spirit, [4] combines deep and shallow convolutions within the same network, providing accurate counting of dense objects (e.g. the UCF50 crowd dataset).
+
+In this paper, however, we aim to apply the dilated convolution method of [25], which has shown to be able to incorporate multiscale perspective information without using multiple inputs or a complicated network architecture, as well as the multicolumn approach of [8, 28] to aggregate multiscale information for the counting problem.
+
+## 3. Method 3.1. Dilated Convolutions For Multicolumn Networks
+
+We propose the use of dilated convolutions as an attractive alternative to the architecture of the HydraCNN
+[18], which seems to saturate in performance at 3 or more columns. We refer to our proposed network as the aggregated multicolumn dilated convolution network1, henceforth shortened as the AMDCN. The architecture of the AMDCN is inspired by the multicolumn counting network of [28]. Extracting features from multiple scales is a good idea when attempting to perform perspective-free counting and increasing the convolution kernel size across columns is an efficient method of doing so. However, the number of parameters increases exponentially as larger kernels are used in these columns to extract features at larger scales. Therefore, we propose using dilated convolutions rather than larger kernels.
+
+Dilated convolutions, as discussed in [25], allow for the exponential increase of the receptive field with a linear increase in the number of parameters with respect to each hidden layer.
+
+In a traditional 2D convolution, we define a real valued function F : Z
+2 → R, an input Ωr = [−*r, r*]
+2 ∈ Z
+2, and a filter function k : Ωr → R. In this case, a convolution operation as defined in [25] is given by
+
+$$(F*k)({\bf p})=\sum_{{\bf s}+{\bf t}={\bf p}}F({\bf s})k({\bf t}).\qquad\qquad(1)$$
+
+A dilated convolution is essentially a generalization of the traditional 2D convolution that allows the operation to skip some inputs. This enables an increase in the size of the filter (i.e. the size of the receptive field) without losing resolution. Formally, we define from [25] the dilated convolution as
+
+$$(F*_{l}k)(\mathbf{p})=\sum_{\mathbf{s}+l\mathbf{t}=\mathbf{p}}F(\mathbf{s})k(\mathbf{t})$$
+$${\mathrm{(2)}}$$
+F(s)k(t) (2)
+where l is the index of the current layer of the convolution.
+
+Using dilations to construct the aggregator in combination with the multicolumn idea will allow for the construction of a network with more than just 3 or 4 columns as in [28] and [8], because the aggregator should prevent the saturation of performance with increasing numbers of columns. Therefore the network will be able to extract useful features from more scales. We take advantage of dilations within the columns as well to provide large receptive fields with fewer parameters.
+
+Looking at more scales should allow for more accurate regression of the density map. However, because not all scales will be relevant, we extend the network beyond a simple 1 × 1 convolution after the merged columns. Instead, we construct a second part of the network, the aggregator, which sets our method apart from [28], [8], and other multicolumn networks. This aggregator is another series of dilated convolutions that should appropriately consolidate the multiscale information collected by the columns. This is a capability of dilated convolutions observed by [25].
+
+While papers such as [28] and [8] have shown that multiple columns and dilated columns are useful in extracting multiscale information, we argue in this paper that the simple aggregator module built using dilated convolutions is able to effectively make use multiscale information from multiple columns. We show compelling evidence for these claims in Section 4.5.
+
+The network as shown in Figure 1 contains 5 columns.
+
+Note that dilations allow us to use more columns for counting than [28] or [8]. Each column looks at a larger scale than the previous (the exact dilations can also be seen in Figure 1). There are 32 feature maps for each convolution, and all inputs are zero padded prior to each convolution in order to maintain the same data shape from input to output. That is, an image input to this network will result in a density map of the same dimensions. All activations in the specified network are ReLUs. Our input pixel values are floating point 32 bit values from 0 to 1. We center our inputs at 0 by subtracting the per channel mean from each channel. When training, we use a scaled mean absolute error for our loss function:
+
+$$L={\frac{1}{n}}\sum_{i=1}^{n}\vert{\hat{y}}_{i}-\gamma y_{i}\vert$$
+
+where γ is the scale factor, yˆiis the prediction, yiis the true value, and n is the number of pixels. We use a scaled mean absolute error because the target values are so small that it is numerically unstable to regress to these values. At testing time, when retrieving the output density map from the network, we scale the pixel values by γ
+−1to obtain the correct value. This approach is more numerically stable and avoids having the network learn to output only zeros by weighting the nonzero values highly. For all our datasets, we set γ = 255.
+
+## 3.2. Experiments
+
+We evaluated the performance of dilated convolutions against various counting methods on a variety of common counting datasets: UCF50 crowd data, TRANCOS traffic data [18], UCSD crowd data [5], and WorldExpo crowd data [27]. For each of these data sets, we used labels given by the corresponding density map for each image. An example of this is shown in Figure 2. We have performed experiments on the four different splits of the UCSD data as used in [18] and the split of the UCSD data as used in [28] (which we call the original split). We also evaluated the performance of our network on the TRANCOS traffic dataset [14]. We have also experimented with higher density datasets for crowd counting, namely WorldExpo and UCF.
+
+We have observed that multicolumn dilations produce density maps (and therefore counts) that often have lower loss than those of HydraCNN [18] and [28]. We measure density map regression loss via a scaled mean absolute error loss during training. We compare accuracy of the counts via mean absolute error for the crowd datasets and the GAME
+metric in the TRANCOS dataset as explained in Section 3.2.2. Beyond the comparison to HydraCNN, we will also compare to other recent convolutional counting methods, especially those of [21], [24], and [4] where possible.
+
+For all datasets, we generally use patched input images and ground truth density maps produced by summing a Gaussian of a fixed size (σ) for each object for training.
+
+This size varies from dataset to dataset, but remains constant within a dataset with the exception of cases in which a perspective map is used. This is explained per dataset. All experiments were performed using Keras with the Adam optimizer [10]. The learning rates used are detailed per dataset.
+
+For testing, we also use patches that can either be directly pieced together or overlapped and averaged except in the case of UCF, for which we run our network on the full image.
+
+$$(3)$$
+
+Furthermore, we performed a set of experiments in which we varied the number of columns from 1 to 5 (simply by including or not including the columns as specified in Figure 1, starting with the smallest filter column and adding larger filter columns one by one). Essentially, the network is allowed to extract information at larger and larger scales in addition to the smaller scales as we include each column. We then performed the same set of experiments, varying the number of columns, but with the aggregator module removed. We perform these experiments on the original split of UCSD as specified in Section 3.2.3 and [5], the TRANCOS dataset, and the WorldExpo dataset because these are relatively large and well defined datasets. We limit the number of epochs to 10 for all of these sets of experiments in order to control for the effect of learning time, and also compare all results using MAE for consistency. These experiments are key to determining the efficacy of the aggregator in effectively combining multiscale information and in providing evidence to support the use of multiple columns to extract multiscale information from images. We report the results of these ablation studies in Section 4.5.
+
+## 3.2.1 Ucf50 Crowd Counting
+
+UCF is a particularly challenging crowd counting dataset.
+
+There are only 50 images in the whole dataset and they are all of varying sizes and from different scenes. The number of people also varies between images from less than 100 to the thousands. The average image has on the order of 1000 people. The difficulty is due to the combination of the very low number of images in the dataset and the fact that the images are all of varying scenes, making high quality generalization crucial. Furthermore, perspective effects are particularly noticeable for many images in this dataset. Despite this, there is no perspective information available for this dataset.
+
+We take 1600 random patches of size 150 × 150 for the training. For testing, we do not densely scan the image as in [18] but instead test on the whole image. In order to standardize the image sizes, we pad each image out with zeros until all images are 1024 × 1024. We then suppress output in the regions where we added padding when testing.
+
+This provides a cleaner resulting density map for these large crowds. The ground truth density maps are produced by annotating each object with a Gaussian of σ = 15.
+
+## 3.2.2 Trancos Traffic Counting
+
+TRANCOS is a traffic counting dataset that comes with its own metric [14]. This metric is known as *GAME*, which stands for Grid Average Mean absolute Error. *GAME* splits a given density map into 4 L grids, or subarrays, and obtains a mean absolute error within each grid separately.
+
+The value of L is a parameter chosen by the user. These individual errors are summed to obtain the final error for a particular image. The intuition behind this metric is that it is desirable to penalize a density map whose overall count might match the ground truth, but whose shape does not match the ground truth [14]. More formally, we define
+
+$$G A M E(L)={\frac{1}{N}}\cdot\sum_{n=1}^{N}\left(\sum_{l=1}^{4^{L}}\!\left|e_{n}^{l}-t_{n}^{l}\right|\right)\qquad(4)$$
+
+where N refers to the number of images, L is the level parameter for *GAME*, e l n is the predicted or estimated count in region l of image n and t l n is the ground truth count in region l of image n [14].
+
+For training this dataset, we take 1600 randomly sampled patches of size 80 × 80. For testing this dataset, we take 80 × 80 non-overlapping patches which we can stitch back together into the full-sized 640 × 480 images. We trained the AMDCN network with density maps produced with a Gaussian of σ = 15 as specified in [18].
+
+## 3.2.3 Ucsd Crowd Counting
+
+The UCSD crowd counting dataset consists of frames of video of a sidewalk. There are relatively few people in view at any given time (approximately 25 on average). Furthermore, because the dataset comes from a video, there are many nearly identical images in the dataset. For this dataset, there have been two different ways to split the data into train and test sets. Therefore, we report results using both methods of splitting the data. The first method consists of four different splits: maximal, downscale, upscale, and minimal.
+
+Minimal is particularly challenging as the train set contains only 10 images. Moreover, upscale appears to be the easiest for the majority of methods [18]. The second method of splitting this data is much more succinct, leaving 1200 images in the testing set and 800 images in the training set [28]. This split comes from the original paper, so we call it the original split [5].
+
+For this dataset, each object is annotated with a 2D Gaussian of covariance Σ = 8 · 12×2. The ground truth map is produced by summing these. When we make use of the perspective maps provided, we divide Σ by the perspective map value at that pixel x, represented by M(x). The provided perspective map for UCSD contains both a horizontal and vertical direction so we take the square root of the provided combined value. For training, we take 1600 random 79 × 119 pixel patches and for testing, we split each test image up into quadrants (which have dimension 79 × 119).
+
+There are two different ways to split the dataset into training and testing sets. We have experimented on the split that gave [18] the best results as well as the split used in [28].
+
+First, we split the dataset into four separate groups of training and testing sets as used in [18] and originally defined by [20]. These groups are "upscale," "maximal,"
+"minimal," and "downscale." We see in Table 3 that the
+"upscale" split and "downscale" split give us state of the art results on counting for this dataset. For this experiment, we sampled 1600 random patches of size 119 × 79 pixels
+(width and height respectively) for the training set and split the test set images into 119 × 79 quadrants that could be reconstructed by piecing them together without overlap. We also added left-right flips of each image to our training data.
+
+We then evaluate the original split. For this experiment, we similarly sampled 1600 random patches of size 119×79 pixels (width and height respectively) for the training set and split the test set images into 119 × 79 quadrants that could be reconstructed by piecing them together without overlap.
+
+## 3.2.4 Worldexpo '10 Crowd Counting
+
+The WorldExpo dataset [27] contains a larger number of people (approximately 50 on average, which is double that of UCSD) and contains images from multiple locations.
+
+Perspective effects are also much more noticeable in this dataset as compared to UCSD. These qualities of the dataset serve to increase the difficulty of counting. Like UCSD, the WorldExpo dataset was constructed from frames of video recordings of crowds. This means that, unlike UCF, this dataset contains a relatively large number of training and testing images. We experiment on this dataset with and without perspective information.
+
+Without perspective maps, we generate label density maps for this dataset in the same manner as previously described: a 2D Gaussian with σ = 15. We take 16000 150 × 150 randomly sampled patches for training. For testing, we densely scan the image, producing 150 × 150 patches at a stride of 100.
+
+When perspective maps are used, however, we follow the procedure as described in [27], which involves estimating a
+"crowd density distribution kernel" as the sum of two 2D
+Gaussians: a symmetric Gaussian for the head and an ellipsoid Gaussian for the body. These are scaled by the perspective map M provided, where M(x) gives the number of pixels that represents a meter at pixel x [27]. Note that the meaning of this perspective map is distinct from the meaning of the perspective map provided for the UCSD dataset.
+
+Using this information, the density contribution from a person with head pixel x is given by the following sum of normalized Gaussians:
+
+$$D_{\bf x}=\frac{1}{||Z||}({\cal N}_{h}({\bf x},\sigma_{h})+{\cal N}_{b}({\bf x}_{b},\Sigma_{b}))\qquad\qquad(5)$$
+
+where xb is the center of the body, which is 0.875 meters down from the head on average, and can be determined from the perspective map M and the head center x [27]. We sum these Gaussians for each person to pro-
+Table 1. Mean absolute error of various methods on UCF crowds
+
+| Method       | MAE    |
+|--------------|--------|
+| AMDCN        | 290.82 |
+| Hydra2s [18] | 333.73 |
+| MCNN [28]    | 377.60 |
+| [27]         | 467.00 |
+| [23]         | 295.80 |
+| [3]          | 318.10 |
+
+duce the final density map. We set σ = 0.2M(x) for Nh and σx = 0.2M(x), σy = 0.5M(x) for Σb in Nb.
+
+## 4. Results 4.1. Ucf Crowd Counting
+
+The UCF dataset is particularly challenging due to the large number of people in the images, the variety of the scenes, as well as the low number of training images. We see in Figure 2 that because the UCF dataset has over 1000 people on average in each image, the shapes output by the network in the density map are not as well defined or separated as in the UCSD dataset.
+
+We report a state of the art result on this dataset in Table 1, following the standard protocol of 5-fold cross validation. Our MAE on the dataset is 290.82, which is approximately 5 lower than the previous state of the art, HydraCNN [18]. This is particularly indicative of the power of an aggregated multicolumn dilation network. Despite not making use of perspective information, the AMDCN is still able to produce highly accurate density maps for UCF.
+
+## 4.2. Trancos Traffic Counting
+
+Our network performs very well on the TRANCOS
+dataset. Indeed, as confirmed by the GAME score, AMDCN produces the most accurate count and shape combined as compared to other methods. Table 2 shows that we achieve state of the art results as measured by the *GAME*
+metric [14] across all levels.
+
+## 4.3. Ucsd Crowd Counting
+
+Results are shown in Table 3 and Figure 3. We see that the "original" split as defined by the creators of the dataset in [5] and used in [28] gives us somewhat worse results for counting on this dataset. Results were consistent over multiple trainings. Again, including the perspective map does not seem to increase performance on this dataset. Despite this, we see in Table 3 and Figure 3 that the results are comparable to the state of the art. In fact, for two of the splits, our proposed network beats the state of the art. For the upscale split, the AMDCN is the state of the art by a large relative margin. This is compelling because it shows that accurate perspective-free counting can be achieved without
+
+| Method                                        | GAME   | GAME   | GAME   | GAME   |       |       |
+|-----------------------------------------------|--------|--------|--------|--------|-------|-------|
+| (L=0)                                         | (L=1)  | (L=2)  | (L=3)  |        |       |       |
+| AMDCN                                         | 9.77   | 13.16  | 15.00  | 15.87  |       |       |
+| [18]                                          | 10.99  | 13.75  | 16.69  | 19.32  |       |       |
+| [15]                                          | +      | SIFT   | 13.76  | 16.72  | 20.72 | 24.36 |
+| from [14] [13] + RGB Norm + Filters from [14] | 17.68  | 19.97  | 23.54  | 25.84  |       |       |
+| HOG-2                                         | 13.29  | 18.05  | 23.65  | 28.41  |       |       |
+| from [14]                                     |        |        |        |        |       |       |
+
+creating image pyramids or requiring perspective maps as labels using the techniques presented by the AMDCN.
+
+## 4.4. Worldexpo '10 Crowd Counting
+
+Our network performs reasonably well on the more challenging WorldExpo dataset. While it does not beat the state of the art, our results are comparable. What is more, we do not need to use the perspective maps to obtain these results.
+
+As seen in Table 4, the AMDCN is capable of incorporating the perspective effects without scaling the Gaussians with perspective information. This shows that it is possible to achieve counting results that approach the state of the art with much simpler labels for the counting training data.
+
+## 4.5. Ablation Studies
+
+We report the results of the ablation studies in Figure 4. We note from these plots that while there is variation in performance, a few trends stand out. Most importantly, the lowest errors are consistently with a combination of a larger number of columns and including the aggregator module.
+
+Notably for the TRANCOS dataset, including the aggregator consistently improves performance. Generally, the aggregator tends to decrease the variance in performance of the network. Some of the variance that we see in the plots can be explained by: (1) for lower numbers of columns, including an aggregator is not as likely to help as there is not much separation of multiscale information across columns and (2) for the UCSD dataset, there is less of a perspective effect than TRANCOS and WorldExpo so a simpler network is more likely to perform comparably to a larger network. These results verify the notion that using more columns increases accuracy, and also support our justification for the use of the aggregator module.
+
+![6_image_0.png](6_image_0.png)
+
+![6_image_1.png](6_image_1.png)
+
+| Method                                  | maximal   | downscale   | upscale   | minimal   | original   |
+|-----------------------------------------|-----------|-------------|-----------|-----------|------------|
+| AMDCN (without perspective information) | 1.63      | 1.43        | 0.63      | 1.71      | 1.74       |
+| AMDCN (with perspective information)    | 1.60      | 1.24        | 1.37      | 1.59      | 1.72       |
+| [18] (with perspective information)     | 1.65      | 1.79        | 1.11      | 1.50      | -          |
+| [18] (without perspective information)  | 2.22      | 1.93        | 1.37      | 2.38      | -          |
+| [15]                                    | 1.70      | 1.28        | 1.59      | 2.02      | -          |
+| [13]                                    | 1.70      | 2.16        | 1.61      | 2.20      | -          |
+| [19]                                    | 1.43      | 1.30        | 1.59      | 1.62      | -          |
+| [2]                                     | 1.24      | 1.31        | 1.69      | 1.49      | -          |
+| [27]                                    | 1.70      | 1.26        | 1.59      | 1.52      | 1.60       |
+| [28]                                    | -         | -           | -         | -         | 1.07       |
+| [1, 28]                                 | -         | -           | -         | -         | 2.16       |
+| [7]                                     | -         | -           | -         | -         | 2.25       |
+| [5]                                     | -         | -           | -         | -         | 2.24       |
+| [6]                                     | -         | -           | -         | -         | 2.07       |
+
+## 5. Conclusion 5.1. Summary
+
+We have proposed the use of aggregated multicolumn dilated convolutions, the AMDCN, as an alternative to the HydraCNN [18] or multicolumn CNN [28] for the vision task of counting objects in images. Inspired by the multicolumn approach to multiscale problems, we also employ dilations to increase the receptive field of our columns. We then aggregate this multiscale information using another series of dilated convolutions to enable a wide network and detect features at more scales. This method takes advantage of the ability of dilated convolutions to provide exponentially increasing receptive fields. We have performed experiments on the challenging UCF crowd counting dataset, the TRANCOS traffic dataset, multiple splits of the UCSD
+crowd counting dataset, and the WorldExpo crowd counting dataset.
+
+![7_image_0.png](7_image_0.png)
+
+| Method                              | MAE   |
+|-------------------------------------|-------|
+| AMDCN (without perspective information)                                     | 16.6  |
+| AMDCN (with perspective information)                                     | 14.9  |
+| LBP+RR [28] (with perspective information)                                     | 31.0  |
+| MCNN [28] (with perspective information)                                     | 11.6  |
+| [27] (with perspective information) | 12.9  |
+
+We obtain superior or comparable results in most of these datasets. The AMDCN is capable of outperforming these approaches completely especially when perspective information is not provided, as in UCF and TRANCOS. These results show that the AMDCN performs surprisingly well and is also robust to scale effects. Further, our ablation study of removing the aggregator network shows that using more columns and an aggregator provides the best accuracy for counting - especially so when there is no perspective information.
+
+## 5.2. Future Work
+
+In addition to an analysis of performance on counting, a density regressor can also be used to locate objects in the image. As mentioned previously, if the regressor is accurate and precise enough, the resulting density map can be used to locate the objects in the image. We expect that in order to do this, one must regress each object to a single point rather than a region specified by a Gaussian. Perhaps this might be accomplished by applying non-maxima suppression to the final layer activations.
+
+Indeed, the method of applying dilated filters to a multicolumn convolutional network in order to enable extracting features of a large number of scales can be applied to various other dense prediction tasks, such as object segmentation at multiple scales or single image depth map prediction.
+
+Though we have only conducted experiments on counting and used 5 columns, the architecture presented can be extended and adapted to a variety of tasks that require information at multiple scales.
+
+## Acknowledgment
+
+This material is based upon work supported by the National Science Foundation under Grant No. 1359275 and 1659788. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation. Furthermore, we acknowledge Kyle Yee and Sridhama Prakhya for their helpful conversations and insights during the research process.
+
+## References
+
+[1] S. An, W. Liu, and S. Venkatesh. Face recognition using kernel ridge regression. In Computer Vision and
+Pattern Recognition, 2007. CVPR'07. IEEE Conference on, pages 1–7. IEEE, 2007.
+[2] C. Arteta, V. Lempitsky, J. A. Noble, and A. Zisserman. Interactive object counting. In *European Conference on Computer Vision*, pages 504–518. Springer,
+2014.
+[3] D. Babu Sam, S. Surya, and R. Venkatesh Babu.
+Switching convolutional neural network for crowd
+[15] V. Lempitsky and A. Zisserman. Learning to count objects in images. In *Advances in Neural Information*
+Processing Systems, pages 1324–1332, 2010.
+
+[16] G. Lin, C. Shen, A. van den Hengel, and I. Reid. Efficient piecewise training of deep structured models for semantic segmentation. In Proceedings of the IEEE
+Conference on Computer Vision and Pattern Recognition, pages 3194–3203, 2016.
+
+[17] H. Noh, S. Hong, and B. Han. Learning deconvolution network for semantic segmentation. In Proceedings of the IEEE International Conference on Computer Vision, pages 1520–1528, 2015.
+
+[18] D. Onoro-Rubio and R. J. Lopez-Sastre. Towards ´
+perspective-free object counting with deep learning.
+
+In *European Conference on Computer Vision*, pages 615–629. Springer, 2016.
+
+[19] V.-Q. Pham, T. Kozakaya, O. Yamaguchi, and R. Okada. Count forest: Co-voting uncertain number of targets using random forest for crowd density estimation. In Proceedings of the IEEE International Conference on Computer Vision, pages 3253–3261, 2015.
+
+[20] D. Ryan, S. Denman, C. Fookes, and S. Sridharan.
+
+Crowd counting using multiple local features. In Digital Image Computing: Techniques and Applications, 2009. DICTA'09., pages 81–88. IEEE, 2009.
+
+[21] S. Segu´ı, O. Pujol, and J. Vitria. Learning to count with deep object features. In *Proceedings of the IEEE*
+Conference on Computer Vision and Pattern Recognition Workshops, pages 90–96, 2015.
+
+[22] J. Selinummi, O. Yli-Harja, and J. A. Puhakka. Software for quantification of labeled bacteria from digital microscope images by automated image analysis.
+
+Biotechniques, 39(6):859, 2005.
+
+[23] V. A. Sindagi and V. M. Patel. Generating high-quality crowd density maps using contextual pyramid cnns.
+
+In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 1861–1870, 2017.
+
+[24] E. Walach and L. Wolf. Learning to count with cnn boosting. In *European Conference on Computer Vision*, pages 660–676. Springer, 2016.
+
+[25] F. Yu and V. Koltun. Multi-scale context aggregation by dilated convolutions. *arXiv preprint*
+arXiv:1511.07122, 2015.
+
+[26] F. Yu, V. Koltun, and T. Funkhouser. Dilated residual networks. *arXiv preprint arXiv:1705.09914*, 2017.
+
+[27] C. Zhang, H. Li, X. Wang, and X. Yang. Crossscene crowd counting via deep convolutional neural networks. In Proceedings of the IEEE Conference on counting. In *Proceedings of the IEEE Conference*
+on Computer Vision and Pattern Recognition, pages 5744���5752, 2017.
+
+[4] L. Boominathan, S. S. Kruthiventi, and R. V. Babu.
+
+Crowdnet: A deep convolutional network for dense crowd counting. In Proceedings of the 2016 ACM on Multimedia Conference, pages 640–644. ACM, 2016.
+
+[5] A. B. Chan, Z.-S. J. Liang, and N. Vasconcelos. Privacy preserving crowd monitoring: Counting people without people models or tracking. In Computer Vision and Pattern Recognition, 2008. CVPR 2008.
+
+IEEE Conference on, pages 1–7. IEEE, 2008.
+
+[6] K. Chen, S. Gong, T. Xiang, and C. Change Loy. Cumulative attribute space for age and crowd density estimation. In *Proceedings of the IEEE conference on*
+computer vision and pattern recognition, pages 2467–
+2474, 2013.
+
+[7] K. Chen, C. C. Loy, S. Gong, and T. Xiang. Feature mining for localised crowd counting.
+
+[8] L.-C. Chen, G. Papandreou, I. Kokkinos, K. Murphy, and A. L. Yuille. Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs. *IEEE Transactions on Pattern Analysis and Machine Intelligence*, 2017.
+
+[9] L.-C. Chen, Y. Yang, J. Wang, W. Xu, and A. L. Yuille.
+
+Attention to scale: Scale-aware semantic image segmentation. In *Proceedings of the IEEE Conference*
+on Computer Vision and Pattern Recognition, pages 3640–3649, 2016.
+
+[10] F. Chollet et al. Keras. https://github.com/
+fchollet/keras, 2015.
+
+[11] A. Dosovitskiy, P. Fischer, E. Ilg, P. Hausser, C. Hazirbas, V. Golkov, P. van der Smagt, D. Cremers, and T. Brox. Flownet: Learning optical flow with convolutional networks. In *Proceedings of the IEEE International Conference on Computer Vision*, pages 2758–
+2766, 2015.
+
+[12] C. Farabet, C. Couprie, L. Najman, and Y. LeCun. Learning hierarchical features for scene labeling. *IEEE transactions on pattern analysis and machine intelligence*, 35(8):1915–1929, 2013.
+
+[13] L. Fiaschi, U. Kothe, R. Nair, and F. A. Hamprecht. ¨
+Learning to count with regression forest and structured labels. In *Pattern Recognition (ICPR), 2012 21st International Conference on*, pages 2685–2688. IEEE,
+2012.
+
+[14] R. Guerrero-Gomez-Olmedo, B. Torre-Jim ´ enez, S. M. ´
+Lopez-Sastre, Roberto Basc ´ on, and D. O ´ noro Rubio. ˜
+Extremely overlapping vehicle counting. In *Iberian*
+Conference on Pattern Recognition and Image Analysis (IbPRIA), 2015.
+
+Computer Vision and Pattern Recognition, pages 833–
+841, 2015.
+
+[28] Y. Zhang, D. Zhou, S. Chen, S. Gao, and Y. Ma.
+
+Single-image crowd counting via multi-column convolutional neural network. In *Proceedings of the IEEE*
+Conference on Computer Vision and Pattern Recognition, pages 589–597, 2016.

data/examples/nougat/multicolcnn.md

@@ -0,0 +1,245 @@
+# An Aggregated Multicolumn Dilated Convolution Network
+
+for Perspective-Free Counting
+
+Diptodip Deb
+
+Georgia Institute of Technology
+
+[email protected]
+
+Jonathan Ventura
+
+University of Colorado Colorado Springs
+
+[email protected]
+
+###### Abstract
+
+We propose the use of dilated filters to construct an aggregation module in a multicolumn convolutional neural network for perspective-free counting. Counting is a common problem in computer vision (e.g. traffic on the street or pedestrians in a crowd). Modern approaches to the counting problem involve the production of a density map via regression whose integral is equal to the number of objects in the image. However, objects in the image can occur at different scales (e.g. due to perspective effects) which can make it difficult for a learning agent to learn the proper density map. While the use of multiple columns to extract multiscale information from images has been shown before, our approach aggregates the multiscale information gathered by the multicolumn convolutional neural network to improve performance. Our experiments show that our proposed network outperforms the state-of-the-art on many benchmark datasets, and also that using our aggregation module in combination with a higher number of columns is beneficial for multiscale counting.
+
+## 1 Introduction
+
+Learning to count the number of objects in an image is a deceptively difficult problem with many interesting applications, such as surveillance [20], traffic monitoring [14] and medical image analysis [22]. In many of these application areas, the objects to be counted vary widely in appearance, size and shape, and labeled training data is typically sparse. These factors pose a significant computer vision and machine learning challenge.
+
+Lempitsky et al. [15] showed that it is possible to learn to count without learning to explicitly detect and localize individual objects. Instead, they propose learning to predict a density map whose integral over the image equals the number of objects in the image. This approach has been adopted by many later works (Cf. [18, 28]).
+
+However, in many counting problems, such as those counting cells in a microscope image, pedestrians in a crowd, or vehicles in a traffic jam, regressors trained on a single image scale are not reliable [18]. This is due to a variety of challenges including overlap of objects and perspective effects which cause significant variance in object shape, size and appearance.
+
+The most successful recent approaches address this issue by explicitly incorporating multi-scale information in the network [18, 28]. These approaches either combine multiple networks which take input patches of different sizes [18] or combine multiple filtering paths ("columns") which have different size filters [28].
+
+Following on the intuition that multiscale integration is key to achieving good counting performance, we propose to incorporate dilated filters [25] into a multicolumn convolutional neural network design [28]. Dilated filters exponentially increase the network's receptive field without an exponential increase in parameters, allowing for efficient use of multiscale information. Convolutional neural networks with dilated filters have proven to provide competitive performance in image segmentation where multiscale analysis is also critical [25, 26]. By incorporating dilated filters into the multicolumn network design, we greatly increase the ability of the network to selectively aggregate multiscale information, without the need for explicit perspective maps during training and testing. We propose the "aggregated multicolumn dilated convolution network" or AMDCN which uses dilations to aggregate multiscale information. Our extensive experimental evaluation shows that this proposed network outperforms previous methods on many benchmark datasets.
+
+## 2 Related Work
+
+Counting using a supervised regressor to formulate a density map was first shown by [15]. In this paper, Lempitsky et al. show that the minimal annotation of a single dot blurred by a Gaussian kernel produces a sufficient density map to train a network to count. All of the counting methods that we examine as well as the method we use inour paper follow this method of producing a density map via regression. This is particularly advantageous because a sufficiently accurate regressor can also locate the objects in the image via this method. However, the Lempitsky paper ignores the issue of perspective scaling and other scaling issues. The work of [27] introduces CNNs (convolutional neural networks) for the purposes of crowd counting, but performs regression on similarly scaled image patches.
+
+These issues are addressed by the work of [18]. Rubio et al. show that a fully convolutional neural network can be used to produce a supervised regressor that produces density maps as in [15]. They further demonstrate a method dubbed HydraCNN which essentially combines multiple convolutional networks that take in differently scaled image patches in order to incorporate multiscale, global information from the image. The premise of this method is that a single regressor will fail to accurately represent the difference in values of the features of an image caused by perspective shifts (scaling effects) [18].
+
+However, the architectures of both [18] and [27] are not fully convolutional due to requiring multiple image patches and, as discussed in [25], the experiments of [11, 17] and [9, 12, 16] leave it unclear as to whether rescaling patches of the image is truly necessary in order to solve dense prediction problems via convolutional neural networks. Moreover, these approaches seem to saturate in performance at three columns, which means the network is extracting information from fewer scales. The work of [25] proposes the use of dilated convolutions as a simpler alternative that does not require sampling of rescaled image patches to provide global, scale-aware information to the network. A fully convolutional approach to multiscale counting has been proposed by [28], in which a multicolumn convolutional network gathers features of different scales by using convolutions of increasing kernel sizes from column to column instead of scaling image patches. Further, DeepLab has used dilated convolutions in multiple columns to extract scale information for segmentation [8]. We build on these approaches with our aggregator module as described in Section 3.1, which should allow for extracting information from more scales.
+
+It should be noted that other methods of counting exist, including training a network to recognize deep object features via only providing the counts of the objects of interest in an image [21] and using CNNs (convolutional neural networks) along with boosting in order to improve the results
+
+Figure 1: Fully convolutional architecture diagram (not to scale). Arrows show separate columns that all take the same input. At the end of the columns, the feature maps are merged (concatenated) together and passed to another series of dilated convolutions: the aggregator, which can aggregate the multiscale information collected by the columns [25]. The input image is I with C channels. The output single channel density map is D, and integrating over this map (summing the pixels) results in the final count. Initial filter sizes are labeled with brackets or lines. Convolution operations are shown as flat rectangles, feature maps are shown as prisms. The number below each filter represents the dilation rate (1 means no dilation).
+
+of regression for production of density maps [24]. In the same spirit, [4] combines deep and shallow convolutions within the same network, providing accurate counting of dense objects (e.g. the UCF50 crowd dataset).
+
+In this paper, however, we aim to apply the dilated convolution method of [25], which has shown to be able to incorporate multiscale perspective information without using multiple inputs or a complicated network architecture, as well as the multicolumn approach of [8, 28] to aggregate multiscale information for the counting problem.
+
+## 3 Method
+
+### Dilated Convolutions for Multicolumn Networks
+
+We propose the use of dilated convolutions as an attractive alternative to the architecture of the HydraCNN [18], which seems to saturate in performance at 3 or more columns. We refer to our proposed network as the aggregated multicolumn dilated convolution network1, henceforth shortened as the AMDCN. The architecture of the AMDCN is inspired by the multicolumn counting network of [28]. Extracting features from multiple scales is a good idea when attempting to perform perspective-free counting and increasing the convolution kernel size across columns is an efficient method of doing so. However, the number of parameters increases exponentially as larger kernels are used in these columns to extract features at larger scales. Therefore, we propose using dilated convolutions rather than larger kernels.
+
+Footnote 1: Implementation available on [https://github.com/dipotdip/counting](https://github.com/dipotdip/counting).
+
+Dilated convolutions, as discussed in [25], allow for the exponential increase of the receptive field with a linear increase in the number of parameters with respect to each hidden layer.
+
+In a traditional 2D convolution, we define a real valued function \(F:\mathbb{Z}^{2}\rightarrow\mathbb{R}\), an input \(\Omega_{r}=[-r,r]^{2}\in\mathbb{Z}^{2}\), and a filter function \(k:\Omega_{r}\rightarrow\mathbb{R}\). In this case, a convolution operation as defined in [25] is given by
+
+\[(F*k)(\mathbf{p})=\sum_{\mathbf{s}+\mathbf{t}=\mathbf{p}}F(\mathbf{s})k( \mathbf{t}). \tag{1}\]
+
+A dilated convolution is essentially a generalization of the traditional 2D convolution that allows the operation to skip some inputs. This enables an increase in the size of the filter (i.e. the size of the receptive field) without losing resolution. Formally, we define from [25] the dilated convolution as
+
+\[(F*_{l}k)(\mathbf{p})=\sum_{\mathbf{s}+l\mathbf{t}=\mathbf{p}}F(\mathbf{s})k( \mathbf{t}) \tag{2}\]
+
+where \(l\) is the index of the current layer of the convolution.
+
+Using dilations to construct the aggregator in combination with the multicolumn idea will allow for the construction of a network with more than just 3 or 4 columns as in [28] and [8], because the aggregator should prevent the saturation of performance with increasing numbers of columns. Therefore the network will be able to extract useful features from more scales. We take advantage of dilations within the columns as well to provide large receptive fields with fewer parameters.
+
+Looking at more scales should allow for more accurate regression of the density map. However, because not all scales will be relevant, we extend the network beyond a simple \(1\times 1\) convolution after the merged columns. Instead, we construct a second part of the network, the aggregator, which sets our method apart from [28, 8], and other multicolumn networks. This aggregator is another series of dilated convolutions that should appropriately consolidate the multiscale information collected by the columns. This is a capability of dilated convolutions observed by [25]. While papers such as [28] and [8] have shown that multiple columns and dilated columns are useful in extracting multi-scale information, we argue in this paper that the simple aggregator module built using dilated convolutions is able to effectively make use multiscale information from multiple columns. We show compelling evidence for these claims in Section 4.5.
+
+The network as shown in Figure 1 contains 5 columns. Note that dilations allow us to use more columns for counting than [28] or [8]. Each column looks at a larger scale than the previous (the exact dilations can also be seen in Figure 1). There are 32 feature maps for each convolution, and all inputs are zero padded prior to each convolution in order to maintain the same data shape from input to output. That is, an image input to this network will result in a density map of the same dimensions. All activations in the specified network are ReLUs. Our input pixel values are floating point 32 bit values from 0 to 1. We center our inputs at 0 by subtracting the per channel mean from each channel. When
+
+Figure 2: UCF sample results. Left: input counting image. Middle: Ground truth density map. Right: AMDCN prediction of density map on test image. The network never saw these images during training. All density maps are one channel only (i.e. grayscale), but are colored here for clarity.
+
+training, we use a scaled mean absolute error for our loss function:
+
+\[L=\frac{1}{n}\sum_{i=1}^{n}|\hat{y}_{i}-\gamma y_{i}| \tag{3}\]
+
+where \(\gamma\) is the scale factor, \(\hat{y}_{i}\) is the prediction, \(y_{i}\) is the true value, and \(n\) is the number of pixels. We use a scaled mean absolute error because the target values are so small that it is numerically unstable to regress to these values. At testing time, when retrieving the output density map from the network, we scale the pixel values by \(\gamma^{-1}\) to obtain the correct value. This approach is more numerically stable and avoids having the network learn to output only zeros by weighting the nonzero values highly. For all our datasets, we set \(\gamma=255\).
+
+### Experiments
+
+We evaluated the performance of dilated convolutions against various counting methods on a variety of common counting datasets: UCF50 crowd data, TRANCOS traffic data [18], UCSD crowd data [5], and WorldExpo crowd data [27]. For each of these data sets, we used labels given by the corresponding density map for each image. An example of this is shown in Figure 2. We have performed experiments on the four different splits of the UCSD data as used in [18] and the split of the UCSD data as used in [28] (which we call the original split). We also evaluated the performance of our network on the TRANCOS traffic dataset [14]. We have also experimented with higher density datasets for crowd counting, namely WorldExpo and UCF.
+
+We have observed that multicolumn dilations produce density maps (and therefore counts) that often have lower loss than those of HydraCNN [18] and [28]. We measure density map regression loss via a scaled mean absolute error loss during training. We compare accuracy of the counts via mean absolute error for the crowd datasets and the GAME metric in the TRANCOS dataset as explained in Section 3.2.2. Beyond the comparison to HydraCNN, we will also compare to other recent convolutional counting methods, especially those of [21], [24], and [4] where possible.
+
+For all datasets, we generally use patched input images and ground truth density maps produced by summing a Gaussian of a fixed size (\(\sigma\)) for each object for training. This size varies from dataset to dataset, but remains constant within a dataset with the exception of cases in which a perspective map is used. This is explained per dataset. All experiments were performed using Keras with the Adam optimizer [10]. The learning rates used are detailed per dataset. For testing, we also use patches that can either be directly pieced together or overlapped and averaged except in the case of UCF, for which we run our network on the full image.
+
+Furthermore, we performed a set of experiments in which we varied the number of columns from 1 to 5 (simply by including or not including the columns as specified in Figure 1, starting with the smallest filter column and adding larger filter columns one by one). Essentially, the network is allowed to extract information at larger and larger scales in addition to the smaller scales as we include each column. We then performed the same set of experiments, varying the number of columns, but with the aggregator module removed. We perform these experiments on the original split of UCSD as specified in Section 3.2.3 and [5], the TRANCOS dataset, and the WorldExpo dataset because these are relatively large and well defined datasets. We limit the number of epochs to 10 for all of these sets of experiments in order to control for the effect of learning time, and also compare all results using MAE for consistency. These experiments are key to determining the efficacy of the aggregator in effectively combining multiscale information and in providing evidence to support the use of multiple columns to extract multiscale information from images. We report the results of these ablation studies in Section 4.5.
+
+#### 3.2.1 UCF50 Crowd Counting
+
+UCF is a particularly challenging crowd counting dataset. There are only 50 images in the whole dataset and they are all of varying sizes and from different scenes. The number of people also varies between images from less than 100 to the thousands. The average image has on the order of 1000 people. The difficulty is due to the combination of the very low number of images in the dataset and the fact that the images are all of varying scenes, making high quality generalization crucial. Furthermore, perspective effects are particularly noticeable for many images in this dataset. Despite this, there is no perspective information available for this dataset.
+
+We take 1600 random patches of size \(150\times 150\) for the training. For testing, we do not densely scan the image as in [18] but instead test on the whole image. In order to standardize the image sizes, we pad each image out with zeros until all images are \(1024\times 1024\). We then suppress output in the regions where we added padding when testing. This provides a cleaner resulting density map for these large crowds. The ground truth density maps are produced by annotating each object with a Gaussian of \(\sigma=15\).
+
+#### 3.2.2 TRANCOS Traffic Counting
+
+TRANCOS is a traffic counting dataset that comes with its own metric [14]. This metric is known as \(GAME\), which stands for Grid Average Mean absolute Error. \(GAME\) splits a given density map into \(4^{L}\) grids, or subarrays, and obtains a mean absolute error within each grid separately. The value of \(L\) is a parameter chosen by the user. Theseindividual errors are summed to obtain the final error for a particular image. The intuition behind this metric is that it is desirable to penalize a density map whose overall count might match the ground truth, but whose shape does not match the ground truth [14]. More formally, we define
+
+\[GAME(L)=\frac{1}{N}\cdot\sum_{n=1}^{N}\left(\sum_{l=1}^{4^{L}}\lvert e_{n}^{l}-t_{ n}^{l}\rvert\right) \tag{4}\]
+
+where \(N\) refers to the number of images, \(L\) is the level parameter for \(GAME\), \(e_{n}^{l}\) is the predicted or estimated count in region \(l\) of image \(n\) and \(t_{n}^{l}\) is the ground truth count in region \(l\) of image \(n\)[14].
+
+For training this dataset, we take 1600 randomly sampled patches of size \(80\times 80\). For testing this dataset, we take \(80\times 80\) non-overlapping patches which we can stitch back together into the full-sized \(640\times 480\) images. We trained the AMDCN network with density maps produced with a Gaussian of \(\sigma=15\) as specified in [18].
+
+#### 3.2.3 UCSD Crowd Counting
+
+The UCSD crowd counting dataset consists of frames of video of a sidewalk. There are relatively few people in view at any given time (approximately 25 on average). Furthermore, because the dataset comes from a video, there are many nearly identical images in the dataset. For this dataset, there have been two different ways to split the data into train and test sets. Therefore, we report results using both methods of splitting the data. The first method consists of four different splits: maximal, downscale, upscale, and minimal. Minimal is particularly challenging as the train set contains only 10 images. Moreover, upscale appears to be the easiest for the majority of methods [18]. The second method of splitting this data is much more succinct, leaving 1200 images in the testing set and 800 images in the training set [28]. This split comes from the original paper, so we call it the original split [5].
+
+For this dataset, each object is annotated with a 2D Gaussian of covariance \(\Sigma=8\cdot\mathbf{1}_{2\times 2}\). The ground truth map is produced by summing these. When we make use of the perspective maps provided, we divide \(\Sigma\) by the perspective map value at that pixel \(\mathbf{x}\), represented by \(M(\mathbf{x})\). The provided perspective map for UCSD contains both a horizontal and vertical direction so we take the square root of the provided combined value. For training, we take 1600 random \(79\times 119\) pixel patches and for testing, we split each test image up into quadrants (which have dimension \(79\times 119\)). There are two different ways to split the dataset into training and testing sets. We have experimented on the split that gave [18] the best results as well as the split used in [28].
+
+First, we split the dataset into four separate groups of training and testing sets as used in [18] and originally defined by [20]. These groups are "upscale," "maximal," "minimal," and "downscale." We see in Table 3 that the "upscale" split and "downscale" split give us state of the art results on counting for this dataset. For this experiment, we sampled 1600 random patches of size \(119\times 79\) pixels (width and height respectively) for the training set and split the test set images into \(119\times 79\) quadrants that could be reconstructed by piecing them together without overlap. We also added left-right flips of each image to our training data.
+
+We then evaluate the original split. For this experiment, we similarly sampled 1600 random patches of size \(119\times 79\) pixels (width and height respectively) for the training set and split the test set images into \(119\times 79\) quadrants that could be reconstructed by piecing them together without overlap.
+
+#### 3.2.4 WorldExpo '10 Crowd Counting
+
+The WorldExpo dataset [27] contains a larger number of people (approximately 50 on average, which is double that of UCSD) and contains images from multiple locations. Perspective effects are also much more noticeable in this dataset as compared to UCSD. These qualities of the dataset serve to increase the difficulty of counting. Like UCSD, the WorldExpo dataset was constructed from frames of video recordings of crowds. This means that, unlike UCF, this dataset contains a relatively large number of training and testing images. We experiment on this dataset with and without perspective information.
+
+Without perspective maps, we generate label density maps for this dataset in the same manner as previously described: a 2D Gaussian with \(\sigma=15\). We take 16000 \(150\times 150\) randomly sampled patches for training. For testing, we densely scan the image, producing \(150\times 150\) patches at a stride of 100.
+
+When perspective maps are used, however, we follow the procedure as described in [27], which involves estimating a "crowd density distribution kernel" as the sum of two 2D Gaussians: a symmetric Gaussian for the head and an ellipsoid Gaussian for the body. These are scaled by the perspective map \(M\) provided, where \(M(\mathbf{x})\) gives the number of pixels that represents a meter at pixel \(\mathbf{x}\)[27]. Note that the meaning of this perspective map is distinct from the meaning of the perspective map provided for the UCSD dataset. Using this information, the density contribution from a person with head pixel \(\mathbf{x}\) is given by the following sum of normalized Gaussians:
+
+\[D_{\mathbf{x}}=\frac{1}{||Z||}(\mathcal{N}_{h}(\mathbf{x},\sigma_{h})+\mathcal{ N}_{b}(\mathbf{x}_{b},\Sigma_{b})) \tag{5}\]
+
+where \(\mathbf{x}_{b}\) is the center of the body, which is 0.875 meters down from the head on average, and can be determined from the perspective map \(M\) and the head center \(\mathbf{x}\)[27]. We sum these Gaussians for each person to pro duce the final density map. We set \(\sigma=0.2M(\mathbf{x})\) for \(\mathcal{N}_{h}\) and \(\sigma_{x}=0.2M(\mathbf{x}),\sigma_{y}=0.5M(\mathbf{x})\) for \(\Sigma_{b}\) in \(\mathcal{N}_{b}\).
+
+## 4 Results
+
+### UCF Crowd Counting
+
+The UCF dataset is particularly challenging due to the large number of people in the images, the variety of the scenes, as well as the low number of training images. We see in Figure 2 that because the UCF dataset has over 1000 people on average in each image, the shapes output by the network in the density map are not as well defined or separated as in the UCSD dataset.
+
+We report a state of the art result on this dataset in Table 1, following the standard protocol of 5-fold cross validation. Our MAE on the dataset is 290.82, which is approximately 5 lower than the previous state of the art, HydraCNN [18]. This is particularly indicative of the power of an aggregated multicolumn dilation network. Despite not making use of perspective information, the AMDCN is still able to produce highly accurate density maps for UCF.
+
+### TranCOS Traffic Counting
+
+Our network performs very well on the TRANCOS dataset. Indeed, as confirmed by the GAME score, AMDCN produces the most accurate count and shape combined as compared to other methods. Table 2 shows that we achieve state of the art results as measured by the \(GAME\) metric [14] across all levels.
+
+### UCSD Crowd Counting
+
+Results are shown in Table 3 and Figure 3. We see that the "original" split as defined by the creators of the dataset in [5] and used in [28] gives us somewhat worse results for counting on this dataset. Results were consistent over multiple trainings. Again, including the perspective map does not seem to increase performance on this dataset. Despite this, we see in Table 3 and Figure 3 that the results are comparable to the state of the art. In fact, for two of the splits, our proposed network beats the state of the art. For the up-scale split, the AMDCN is the state of the art by a large relative margin. This is compelling because it shows that accurate perspective-free counting can be achieved without creating image pyramids or requiring perspective maps as labels using the techniques presented by the AMDCN.
+
+### WorldExpo '10 Crowd Counting
+
+Our network performs reasonably well on the more challenging WorldExpo dataset. While it does not beat the state of the art, our results are comparable. What is more, we do not need to use the perspective maps to obtain these results. As seen in Table 4, the AMDCN is capable of incorporating the perspective effects without scaling the Gaussians with perspective information. This shows that it is possible to achieve counting results that approach the state of the art with much simpler labels for the counting training data.
+
+### Ablation Studies
+
+We report the results of the ablation studies in Figure 4. We note from these plots that while there is variation in performance, a few trends stand out. Most importantly, the lowest errors are consistently with a combination of a larger number of columns and including the aggregator module. Notably for the TRANCOS dataset, including the aggregator consistently improves performance. Generally, the aggregator tends to decrease the variance in performance of the network. Some of the variance that we see in the plots can be explained by: (1) for lower numbers of columns, including an aggregator is not as likely to help as there is not much separation of multiscale information across columns and (2) for the UCSD dataset, there is less of a perspective effect than TRANCOS and WorldExpo so a simpler network is more likely to perform comparably to a larger network. These results verify the notion that using more columns increases accuracy, and also support our justification for the use of the aggregator module.
+
+\begin{table}
+\begin{tabular}{|l|l|} \hline
+**Method** & **MAE** \\ \hline AMDCN & **290.82** \\ \hline Hydra2s [18] & 333.73 \\ \hline MCNN [28] & 377.60 \\ \hline [27] & 467.00 \\ \hline [23] & 295.80 \\ \hline [3] & 318.10 \\ \hline \end{tabular}
+\end{table}
+Table 1: Mean absolute error of various methods on UCF crowds
+
+\begin{table}
+\begin{tabular}{|c|l|l|l|l|} \hline
+**Method** & \begin{tabular}{l} **GAME** \\ **(L=0)** \\ \end{tabular} & \begin{tabular}{l} **GAME** \\ **(L=1)** \\ \end{tabular} & \begin{tabular}{l} **GAME** \\ **(L=2)** \\ \end{tabular} & 
+\begin{tabular}{l} **GAME** \\ **(L=3)** \\ \end{tabular} \\ \hline AMDCN & **9.77** & **13.16** & **15.00** & **15.87** \\ \hline [18] & 10.99 & 13.75 & 16.69 & 19.32 \\ \hline [15] + SIFT from [14] & 13.76 & 16.72 & 20.72 & 24.36 \\ \hline [13] + RGB Norm + Filters from [14] & 17.68 & 19.97 & 23.54 & 25.84 \\ \hline HOG-2 from [14] & 13.29 & 18.05 & 23.65 & 28.41 \\ \hline \end{tabular}
+\end{table}
+Table 2: Mean absolute error of various methods on TRANCOS traffic 
+
+## 5 Conclusion
+
+### Summary
+
+We have proposed the use of aggregated multicolumn dilated convolutions, the AMDCN, as an alternative to the HydraCNN [18] or multicolumn CNN [28] for the vision task of counting objects in images. Inspired by the multicolumn approach to multiscale problems, we also employ dilations to increase the receptive field of our columns. We then aggregate this multiscale information using another series of dilated convolutions to enable a wide network and detect features at more scales. This method takes advantage of the ability of dilated convolutions to provide exponentially increasing receptive fields. We have performed experiments on the challenging UCF crowd counting dataset, the TRANCOS traffic dataset, multiple splits of the UCSD crowd counting dataset, and the WorldExpo crowd counting dataset.
+
+\begin{table}
+\begin{tabular}{|l|l|l|l|l|l|} \hline
+**Method** & **maximal** & **downscale** & **upscale** & **minimal** & **original** \\ \hline AMDCN (**without perspective information**) & 1.63 & 1.43 & **0.63** & 1.71 & 1.74 \\ \hline AMDCN (with perspective information) & 1.60 & **1.24** & 1.37 & 1.59 & 1.72 \\ \hline
+[18] (with perspective information) & 1.65 & 1.79 & 1.11 & 1.50 & - \\ \hline
+[18] (without perspective information) & 2.22 & 1.93 & 1.37 & 2.38 & - \\ \hline
+[15] & 1.70 & 1.28 & 1.59 & 2.02 & - \\ \hline
+[13] & 1.70 & 2.16 & 1.61 & 2.20 & - \\ \hline
+[19] & 1.43 & 1.30 & 1.59 & 1.62 & - \\ \hline
+[2] & **1.24** & 1.31 & 1.69 & **1.49** & - \\ \hline
+[27] & 1.70 & 1.26 & 1.59 & 1.52 & 1.60 \\ \hline
+[28] & - & - & - & - & **1.07** \\ \hline
+[1, 28] & - & - & - & - & 2.16 \\ \hline
+[7] & - & - & - & - & 2.25 \\ \hline
+[5] & - & - & - & - & 2.24 \\ \hline
+[6] & - & - & - & - & 2.07 \\ \hline \end{tabular}
+\end{table}
+Table 3: Mean absolute error of various methods on UCSD crowds
+
+Figure 3: UCSD crowd counting dataset. Both plots show comparisons of predicted and ground truth counts over time. While AMDCN does not beat the state of the art on the original split, the predictions still follow the true counts reasonably. The jump in the original split is due to that testing set including multiple scenes of highly varying counts.
+
+We obtain superior or comparable results in most of these datasets. The AMDCN is capable of outperforming these approaches completely especially when perspective information is not provided, as in UCF and TRANCOS. These results show that the AMDCN performs surprisingly well and is also robust to scale effects. Further, our ablation study of removing the aggregator network shows that using more columns and an aggregator provides the best accuracy for counting -- especially so when there is no perspective information.
+
+### Future Work
+
+In addition to an analysis of performance on counting, a density regressor can also be used to locate objects in the image. As mentioned previously, if the regressor is accurate and precise enough, the resulting density map can be used to locate the objects in the image. We expect that in order to do this, one must regress each object to a single point rather than a region specified by a Gaussian. Perhaps this might be accomplished by applying non-maxima suppression to the final layer activations.
+
+Indeed, the method of applying dilated filters to a multi-column convolutional network in order to enable extracting features of a large number of scales can be applied to various other dense prediction tasks, such as object segmentation at multiple scales or single image depth map prediction. Though we have only conducted experiments on counting and used 5 columns, the architecture presented can be extended and adapted to a variety of tasks that require information at multiple scales.
+
+## Acknowledgment
+
+This material is based upon work supported by the National Science Foundation under Grant No. 1359275 and 1659788. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation. Furthermore, we acknowledge Kyle Yee and Sridhama Prakhya for their helpful conversations and insights during the research process.
+
+## References
+
+* [1] S. An, W. Liu, and S. Venkatesh. Face recognition using kernel ridge regression. In _Computer Vision and Pattern Recognition, 2007. CVPR'07. IEEE Conference on_, pages 1-7. IEEE, 2007.
+* [2] C. Arteta, V. Lempitsky, J. A. Noble, and A. Zisserman. Interactive object counting. In _European Conference on Computer Vision_, pages 504-518. Springer, 2014.
+* [3] D. Babu Sam, S. Surya, and R. Venkatesh Babu. Switching convolutional neural network for crowd
+
+\begin{table}
+\begin{tabular}{|l|c|} \hline
+**Method** & **MAE** \\ \hline AMDCN **(without perspective information)** & 16.6 \\ \hline AMDCN (with perspective information) & 14.9 \\ \hline LBP+RR [28] (with perspective information) & 31.0 \\ \hline MCNN [28] (with perspective information) & **11.6** \\ \hline
+[27] (with perspective information) & 12.9 \\ \hline \end{tabular}
+\end{table}
+Table 4: Mean absolute error of various methods on WorldExpo crowds
+
+Figure 4: Ablation studies on various datasets in which the number of columns is varied and the aggregator is included or not included. The results generally support the use of more columns and an aggregator module.
+
+counting. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_, pages 5744-5752, 2017.
+* [4] L. Boominathan, S. S. Kruthiventi, and R. V. Babu. Crowdnet: A deep convolutional network for dense crowd counting. In _Proceedings of the 2016 ACM on Multimedia Conference_, pages 640-644. ACM, 2016.
+* [5] A. B. Chan, Z.-S. J. Liang, and N. Vasconcelos. Privacy preserving crowd monitoring: Counting people without people models or tracking. In _Computer Vision and Pattern Recognition, 2008. CVPR 2008. IEEE Conference on_, pages 1-7. IEEE, 2008.
+* [6] K. Chen, S. Gong, T. Xiang, and C. Change Loy. Cumulative attribute space for age and crowd density estimation. In _Proceedings of the IEEE conference on computer vision and pattern recognition_, pages 2467-2474, 2013.
+* [7] K. Chen, C. C. Loy, S. Gong, and T. Xiang. Feature mining for localised crowd counting.
+* [8] L.-C. Chen, G. Papandreou, I. Kokkinos, K. Murphy, and A. L. Yuille. Deeplab: Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs. _IEEE Transactions on Pattern Analysis and Machine Intelligence_, 2017.
+* [9] L.-C. Chen, Y. Yang, J. Wang, W. Xu, and A. L. Yuille. Attention to scale: Scale-aware semantic image segmentation. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_, pages 3640-3649, 2016.
+* [10] F. Chollet et al. Keras. [https://github.com/fchollet/keras](https://github.com/fchollet/keras), 2015.
+* [11] A. Dosovitskiy, P. Fischer, E. Ilg, P. Hausser, C. Hazirbas, V. Golkov, P. van der Smagt, D. Cremers, and T. Brox. Flownet: Learning optical flow with convolutional networks. In _Proceedings of the IEEE International Conference on Computer Vision_, pages 2758-2766, 2015.
+* [12] C. Farabet, C. Couprie, L. Najman, and Y. LeCun. Learning hierarchical features for scene labeling. _IEEE transactions on pattern analysis and machine intelligence_, 35(8):1915-1929, 2013.
+* [13] L. Fiaschi, U. Kothe, R. Nair, and F. A. Hamprecht. Learning to count with regression forest and structured labels. In _Pattern Recognition (ICPR), 2012 21st International Conference on_, pages 2685-2688. IEEE, 2012.
+* [14] R. Guerrero-Gomez-Olmedo, B. Torre-Jimenez, S. M. Lopez-Sastre, Roberto Bascon, and D. Onoro Rubio. Extremely overlapping vehicle counting. In _Iberian Conference on Pattern Recognition and Image Analysis (IbPRIA)_, 2015.
+* [15] V. Lempitsky and A. Zisserman. Learning to count objects in images. In _Advances in Neural Information Processing Systems_, pages 1324-1332, 2010.
+* [16] G. Lin, C. Shen, A. van den Hengel, and I. Reid. Efficient piecewise training of deep structured models for semantic segmentation. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_, pages 3194-3203, 2016.
+* [17] H. Noh, S. Hong, and B. Han. Learning deconvolution network for semantic segmentation. In _Proceedings of the IEEE International Conference on Computer Vision_, pages 1520-1528, 2015.
+* [18] D. Onoro-Rubio and R. J. Lopez-Sastre. Towards perspective-free object counting with deep learning. In _European Conference on Computer Vision_, pages 615-629. Springer, 2016.
+* [19] V.-Q. Pham, T. Kozakaya, O. Yamaguchi, and R. Okada. Count forest: Co-voting uncertain number of targets using random forest for crowd density estimation. In _Proceedings of the IEEE International Conference on Computer Vision_, pages 3253-3261, 2015.
+* [20] D. Ryan, S. Denman, C. Fookes, and S. Sridharan. Crowd counting using multiple local features. In _Digital Image Computing: Techniques and Applications, 2009. DICTA'09._, pages 81-88. IEEE, 2009.
+* [21] S. Segui, O. Pujol, and J. Vitria. Learning to count with deep object features. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops_, pages 90-96, 2015.
+* [22] J. Selinummi, O. Yli-Harja, and J. A. Puhakka. Software for quantification of labeled bacteria from digital microscope images by automated image analysis. _Biotechniques_, 39(6):859, 2005.
+* [23] V. A. Sindagi and V. M. Patel. Generating high-quality crowd density maps using contextual pyramid cnns. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_, pages 1861-1870, 2017.
+* [24] E. Walach and L. Wolf. Learning to count with cnn boosting. In _European Conference on Computer Vision_, pages 660-676. Springer, 2016.
+* [25] F. Yu and V. Koltun. Multi-scale context aggregation by dilated convolutions. _arXiv preprint arXiv:1511.07122_, 2015.
+* [26] F. Yu, V. Koltun, and T. Funkhouser. Dilated residual networks. _arXiv preprint arXiv:1705.09914_, 2017.
+* [27] C. Zhang, H. Li, X. Wang, and X. Yang. Cross-scene crowd counting via deep convolutional neural networks. In _Proceedings of the IEEE Conference on * [28] Y. Zhang, D. Zhou, S. Chen, S. Gao, and Y. Ma. Single-image crowd counting via multi-column convolutional neural network. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_, pages 589-597, 2016.

data/examples/nougat/switch_transformers.md

@@ -0,0 +1,528 @@
+# Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
+
+William Fedus
+
+1. JAX code for Switch Transformer and all model checkpoints are available at [https://github.com/google-research/t5x](https://github.com/google-research/t5x)
+
+1. [email protected]
+
+Barret Zoph
+
+[email protected]
+
+Noam Shazeer
+
+[email protected]
+
+Google, Mountain View, CA 94043, USA
+
+###### Abstract
+
+In deep learning, models typically reuse the same parameters for all inputs. Mixture of Experts (MoE) models defy this and instead select _different_ parameters for each incoming example. The result is a sparsely-activated model--with an outrageous number of parameters--but a constant computational cost. However, despite several notable successes of MoE, widespread adoption has been hindered by complexity, communication costs, and training instability. We address these with the introduction of the Switch Transformer. We simplify the MoE routing algorithm and design intuitive improved models with reduced communication and computational costs. Our proposed training techniques mitigate the instabilities, and we show large sparse models may be trained, for the first time, with lower precision (bfloat16) formats. We design models based off T5-Base and T5-Large (Raffel et al., 2019) to obtain up to 7x increases in pre-training speed with the same computational resources. These improvements extend into multilingual settings where we measure gains over the mT5-Base version across all 101 languages. Finally, we advance the current scale of language models by pre-training up to trillion parameter models on the "Colossal Clean Crawled Corpus", and achieve a 4x speedup over the T5-XXL model.12
+
+Footnote 1: License: CC-BY 4.0, see [https://creativecommons.org/licenses/by/4.0/](https://creativecommons.org/licenses/by/4.0/). Attribution requirements are provided at [http://jmlr.org/papers/v23/21-0998.html](http://jmlr.org/papers/v23/21-0998.html).
+
+ mixture-of-experts, natural language processing, sparsity, large-scale machine learning, distributed computing
+###### Contents
+
+* 1 Introduction
+* 2 Switch Transformer
+	* 2.1 Simplifying Sparse Routing
+	* 2.2 Efficient Sparse Routing
+	* 2.3 Putting It All Together: The Switch Transformer
+	* 2.4 Improved Training and Fine-Tuning Techniques
+* 3 Scaling Properties
+	* 3.1 Scaling Results on a Step-Basis
+	* 3.2 Scaling Results on a Time-Basis
+	* 3.3 Scaling Versus a Larger Dense Model
+* 4 Downstream Results
+	* 4.1 Fine-Tuning
+	* 4.2 Distillation
+	* 4.3 Multilingual Learning
+* 5 Designing Models with Data, Model, and Expert-Parallelism
+	* 5.1 Data Parallelism
+	* 5.2 Model Parallelism
+	* 5.3 Model and Data Parallelism
+	* 5.4 Expert and Data Parallelism
+	* 5.5 Expert, Model and Data Parallelism
+	* 5.6 Towards Trillion Parameter Models
+* 6 Related Work
+* 7 Discussion
+* 8 Future Work
+* 9 Conclusion
+* A Switch for Attention
+* B Preventing Token Dropping with _No-Token-Left-Behind_
+* C Encouraging Exploration Across Experts
+* D Switch Transformers in Lower Compute Regimes
+* E Relation of Upstream to Downstream Model Performance
+* F Pseudo Code for Switch Transformers
+
+## 1 Introduction
+
+Large scale training has been an effective path towards flexible and powerful neural language models (Radford et al., 2018; Kaplan et al., 2020; Brown et al., 2020). Simple architectures--backed by a generous computational budget, data set size and parameter count--surpass more complicated algorithms (Sutton, 2019). An approach followed in Radford et al. (2018); Raffel et al. (2019); Brown et al. (2020) expands the model size of a densely-activated Transformer (Vaswani et al., 2017). While effective, it is also extremely computationally intensive (Strubell et al., 2019). Inspired by the success of model scale, but seeking greater computational efficiency, we instead propose a _sparsely-activated_ expert model: the Switch Transformer. In our case the sparsity comes from activating a _subset_ of the neural network weights for each incoming example.
+
+Sparse training is an active area of research and engineering (Gray et al., 2017; Gale et al., 2020), but as of today, machine learning libraries and hardware accelerators still cater to dense matrix multiplications. To have an efficient sparse algorithm, we start with the Mixture-of-Expert (MoE) paradigm (Jacobs et al., 1991; Jordan and Jacobs, 1994; Shazeer et al., 2017), and simplify it to yield training stability and computational benefits. MoE models have had notable successes in machine translation (Shazeer et al., 2017, 2018; Lepikhin et al., 2020), however, widespread adoption is hindered by complexity, communication costs, and training instabilities.
+
+We address these issues, and then go beyond translation, to find that these class of algorithms are broadly valuable in natural language. We measure superior scaling on a diverse set of natural language tasks and across three regimes in NLP: pre-training, fine-tuning and multi-task training. While this work focuses on scale, we also show that the Switch Transformer architecture not only excels in the domain of supercomputers, but is
+
+Figure 1: Scaling and sample efficiency of Switch Transformers. Left Plot: Scaling properties for increasingly sparse (more experts) Switch Transformers. Right Plot: Negative log perplexity comparing Switch Transformers to T5 (Raffel et al., 2019) models using the same compute budget.
+
+beneficial even with only a few computational cores. Further, our large sparse models can be distilled (Hinton et al., 2015) into small dense versions while preserving 30% of the sparse model quality gain. Our contributions are the following:
+
+* The Switch Transformer architecture, which simplifies and improves over Mixture of Experts.
+* Scaling properties and a benchmark against the strongly tuned T5 model (Raffel et al., 2019) where we measure 7x+ pre-training speedups while still using the same FLOPS per token. We further show the improvements hold even with limited computational resources, using as few as two experts.
+* Successful distillation of sparse pre-trained and specialized fine-tuned models into small dense models. We reduce the model size by up to 99% while preserving 30% of the quality gains of the large sparse teacher.
+* Improved pre-training and fine-tuning techniques: **(1)** selective precision training that enables training with lower bfloat16 precision **(2)** an initialization scheme that allows for scaling to a larger number of experts and **(3)** increased expert regularization that improves sparse model fine-tuning and multi-task training.
+* A measurement of the pre-training benefits on multilingual data where we find a universal improvement across all 101 languages and with 91% of languages benefiting from 4x+ speedups over the mT5 baseline (Xue et al., 2020).
+* An increase in the scale of neural language models achieved by efficiently combining data, model, and expert-parallelism to create models with up to a trillion parameters. These models improve the pre-training speed of a strongly tuned T5-XXL baseline by 4x.
+
+## 2 Switch Transformer
+
+The guiding design principle for Switch Transformers is to maximize the parameter count of a Transformer model (Vaswani et al., 2017) in a simple and computationally efficient way. The benefit of scale was exhaustively studied in Kaplan et al. (2020) which uncovered power-law scaling with model size, data set size and computational budget. Importantly, this work advocates training large models on relatively small amounts of data as the computationally optimal approach.
+
+Heading these results, we investigate a fourth axis: increase the _parameter count_ while keeping the floating point operations (FLOPs) per example constant. Our hypothesis is that the parameter count, independent of total computation performed, is a separately important axis on which to scale. We achieve this by designing a sparsely activated model that efficiently uses hardware designed for dense matrix multiplications such as GPUs and TPUs. Our work here focuses on TPU architectures, but these class of models may be similarly trained on GPU clusters. In our distributed training setup, our sparsely activated layers split _unique_ weights on different devices. Therefore, the weights of the model increase with the number of devices, all while maintaining a manageable memory and computational footprint on each device.
+
+### Simplifying Sparse Routing
+
+**Mixture of Expert Routing.** Shazeer et al. (2017) proposed a natural language Mixture-of-Experts (MoE) layer which takes as an input a token representation \(x\) and then routes this to the best determined top-\(k\) experts, selected from a set \(\{E_{i}(x)\}_{i=1}^{N}\) of \(N\) experts. The router variable \(W_{r}\) produces logits \(h(x)=W_{r}\cdot x\) which are normalized via a softmax distribution over the available \(N\) experts at that layer. The gate-value for expert \(i\) is given by,
+
+\[p_{i}(x)=\frac{e^{h(x)_{i}}}{\sum_{j}^{N}e^{h(x)_{j}}}. \tag{1}\]
+
+The top-\(k\) gate values are selected for routing the token \(x\). If \(\mathcal{T}\) is the set of selected top-\(k\) indices then the output computation of the layer is the linearly weighted combination of each expert's computation on the token by the gate value,
+
+\[y=\sum_{i\in\mathcal{T}}p_{i}(x)E_{i}(x). \tag{2}\]
+
+**Switch Routing: Rethinking Mixture-of-Experts.** Shazeer et al. (2017) conjectured that routing to \(k>1\) experts was necessary in order to have non-trivial gradients to the routing functions. The authors intuited that learning to route would not work without the ability to compare at least two experts. Ramachandran and Le (2018) went further to
+
+Figure 2: Illustration of a Switch Transformer encoder block. We replace the dense feed forward network (FFN) layer present in the Transformer with a sparse Switch FFN layer (light blue). The layer operates independently on the tokens in the sequence. We diagram two tokens (\(x_{1}=\) “More” and \(x_{2}=\) “Parameters” below) being routed (solid lines) across four FFN experts, where the router independently routes each token. The switch FFN layer returns the output of the selected FFN multiplied by the router gate value (dotted-line).
+
+study the top-\(k\) decision and found that higher \(k\)-values in lower layers in the model were important for models with many routing layers. Contrary to these ideas, we instead use a simplified strategy where we route to only a _single_ expert. We show this simplification preserves model quality, reduces routing computation and performs better. This \(k=1\) routing strategy is later referred to as a Switch layer. Note that for both MoE and Switch Routing, the gate value \(p_{i}(x)\) in Equation 2 permits differentiability of the router.
+
+The benefits for the Switch layer are three-fold: **(1)** The router computation is reduced as we are only routing a token to a single expert. **(2)** The batch size (expert capacity) of each expert can be at least halved since each token is only being routed to a single expert.3
+
+Footnote 3: See Section 2.2 for a technical description.
+
+**(3)** The routing implementation is simplified and communication costs are reduced. Figure 3 shows an example of routing with different expert capacity factors.
+
+### Efficient Sparse Routing
+
+We use Mesh-Tensorflow (MTF) (Shazeer et al., 2018) which is a library, with similar semantics and API to Tensorflow (Abadi et al., 2016) that facilitates efficient distributed data and model parallel architectures. It does so by abstracting the physical set of cores to a logical mesh of processors. Tensors and computations may then be sharded per named dimensions, facilitating easy partitioning of models across dimensions. We design our model with TPUs in mind, which require statically declared sizes. Below we describe our distributed Switch Transformer implementation.
+
+Figure 3: Illustration of token routing dynamics. Each expert processes a fixed batch-size of tokens modulated by the _capacity factor_. Each token is routed to the expert with the highest router probability, but each expert has a fixed batch size of (total_tokens / num_experts) \(\times\) capacity_factor. If the tokens are unevenly dispatched then certain experts will overflow (denoted by dotted red lines), resulting in these tokens not being processed by this layer. A larger capacity factor alleviates this overflow issue, but also increases computation and communication costs (depicted by padded white/empty slots).
+
+**Distributed Switch Implementation.** All of our tensor shapes are statically determined at compilation time, but our computation is _dynamic_ due to the routing decisions at training and inference. Because of this, one important technical consideration is how to set the _expert capacity_. The expert capacity--the number of tokens each expert computes--is set by evenly dividing the number of tokens in the batch across the number of experts, and then further expanding by a _capacity factor_,
+
+\[\text{expert capacity}= \left(\frac{\text{tokens per batch}}{\text{number of experts}}\right) \times\text{capacity factor}. \tag{3}\]
+
+A capacity factor greater than 1.0 creates additional buffer to accommodate for when tokens are not perfectly balanced across experts. If too many tokens are routed to an expert (referred to later as dropped tokens), computation is skipped and the token representation is passed directly to the next layer through the residual connection. Increasing the expert capacity is not without drawbacks, however, since high values will result in wasted computation and memory. This trade-off is explained in Figure 3. Empirically we find ensuring lower rates of dropped tokens are important for the scaling of sparse expert-models. Throughout our experiments we didn't notice any dependency on the number of experts for the number of tokens dropped (typically \(<1\%\)). Using the auxiliary load balancing loss (next section) with a high enough coefficient ensured good load balancing. We study the impact that these design decisions have on model quality and speed in Table 1.
+
+**A Differentiable Load Balancing Loss.** To encourage a balanced load across experts we add an auxiliary loss (Shazeer et al., 2017, 2018; Lepikhin et al., 2020). As in Shazeer et al. (2018); Lepikhin et al. (2020), Switch Transformers simplifies the original design in Shazeer et al. (2017) which had separate load-balancing and importance-weighting losses. For each Switch layer, this auxiliary loss is added to the total model loss during training. Given \(N\) experts indexed by \(i=1\) to \(N\) and a batch \(\mathcal{B}\) with \(T\) tokens, the auxiliary loss is computed as the scaled dot-product between vectors \(f\) and \(P\),
+
+\[\text{loss}=\alpha\cdot N\cdot\sum_{i=1}^{N}f_{i}\cdot P_{i} \tag{4}\]
+
+where \(f_{i}\) is the fraction of tokens dispatched to expert \(i\),
+
+\[f_{i}=\frac{1}{T}\sum_{x\in\mathcal{B}}\mathbbm{1}\{\text{argmax}\:p(x)=i\} \tag{5}\]
+
+and \(P_{i}\) is the fraction of the router probability allocated for expert \(i\), 2
+
+Footnote 2: A potential source of confusion: \(p_{i}(x)\) is the probability of routing token \(x\) to expert \(i\). \(P_{i}\) is the probability fraction to expert \(i\) across _all tokens_ in the batch \(\mathcal{B}\).
+
+\[P_{i}=\frac{1}{T}\sum_{x\in\mathcal{B}}p_{i}(x). \tag{6}\]
+
+Since we seek uniform routing of the batch of tokens across the \(N\) experts, we desire both vectors to have values of \(1/N\). The auxiliary loss of Equation 4 encourages uniform routing since it is minimized under a uniform distribution. The objective can also be differentiated asthe \(P\)-vector is differentiable, but the \(f\)-vector is not. The final loss is multiplied by expert count \(N\) to keep the loss constant as the number of experts varies since under uniform routing \(\sum_{i=1}^{N}(f_{i}\cdot P_{1})=\sum_{i=1}^{N}(\frac{1}{N}\cdot\frac{1}{N})= \frac{1}{N}\). Finally, a hyper-parameter \(\alpha\) is a multiplicative coefficient for these auxiliary losses; throughout this work we use an \(\alpha=10^{-2}\) which was sufficiently large to ensure load balancing while small enough to not to overwhelm the primary cross-entropy objective. We swept hyper-parameter ranges of \(\alpha\) from \(10^{-1}\) to \(10^{-5}\) in powers of 10 and found \(10^{-2}\) balanced load quickly without interfering with training loss.
+
+### Putting It All Together: The Switch Transformer
+
+Our first test of the Switch Transformer starts with pre-training on the "Colossal Clean Crawled Corpus" (C4), introduced in (Raffel et al., 2019). For our pre-training objective, we use a masked language modeling task (Taylor, 1953; Fedus et al., 2018; Devlin et al., 2018) where the model is trained to predict missing tokens. In our pre-training setting, as determined in Raffel et al. (2019) to be optimal, we drop out 15% of tokens and then replace the masked sequence with a single sentinel token. To compare our models, we record the negative log perplexity.4 Throughout all tables in the paper, \(\uparrow\) indicates that a higher value for that metric is better and vice-versa for \(\downarrow\). A comparison of all the models studied in this work are in Table 9.
+
+Footnote 4: We use log base-\(e\) for this metric so the units are nats.
+
+A head-to-head comparison of the Switch Transformer and the MoE Transformer is presented in Table 1. Our Switch Transformer model is FLOP-matched to 'T5-Base' (Raffel et al., 2019) (same amount of computation per token is applied). The MoE Transformer, using top-2 routing, has two experts which each apply a separate FFN to each token and thus its FLOPS are larger. All models were trained for the same number of steps on identical hardware. Note that the MoE model going from capacity factor 2.0 to 1.25 actually slows down (840 to 790) in the above experiment setup, which is unexpected.5
+
+Footnote 5: Note that speed measurements are both a function of the algorithm and the implementation details. Switch Transformer reduces the necessary computation relative to MoE (algorithm), but the final speed differences are impacted by low-level optimizations (implementation).
+
+We highlight three key findings from Table 1: **(1)** Switch Transformers outperform both carefully tuned dense models and MoE Transformers on a speed-quality basis. For a fixed amount of computation and wall-clock time, Switch Transformers achieve the best result. **(2)** The Switch Transformer has a smaller computational footprint than the MoE counterpart. If we increase its size to match the training speed of the MoE Transformer, we find this outperforms all MoE and Dense models on a per step basis as well. **(3)** Switch Transformers perform better at lower capacity factors (1.0, 1.25). Smaller expert capacities are indicative of the scenario in the large model regime where model memory is very scarce and the capacity factor will want to be made as small as possible.
+
+### Improved Training and Fine-Tuning Techniques
+
+Sparse expert models may introduce training difficulties over a vanilla Transformer. Instability can result because of the hard-switching (routing) decisions at each of these layers. Further, low precision formats like bfloat16 (Wang and Kanwar, 2019) can exacerbate issuesin the softmax computation for our router. We describe training difficulties here and the methods we use to overcome them to achieve stable and scalable training.
+
+**Selective precision with large sparse models.** Model instability hinders the ability to train using efficient bfloat16 precision, and as a result, Lepikhin et al. (2020) trains with float32 precision throughout their MoE Transformer. However, we show that by instead _selectively casting_ to float32 precision within a localized part of the model, stability may be achieved, without incurring expensive communication cost of float32 tensors. This technique is inline with modern mixed precision training strategies where certain parts of the model and gradient updates are done in higher precision Micikevicius et al. (2017). Table 2 shows that our approach permits nearly equal speed to bfloat16 training while conferring the training stability of float32.
+
+To achieve this, we cast the router input to float32 precision. The router function takes the tokens as input and produces the dispatch and combine tensors used for the selection and recombination of expert computation (refer to Code Block 15 in the Appendix for details). Importantly, the float32 precision is only used _within_ the body of the router function--on computations local to that device. Because the resulting dispatch and combine tensors are recast to bfloat16 precision at the end of the function, no expensive float32 tensors
+
+\begin{table}
+\begin{tabular}{c c c c c} \hline \hline Model & Capacity & Quality after & Time to Quality & Speed (\(\uparrow\)) \\  & Factor & 100k steps (\(\uparrow\)) & Threshold (\(\downarrow\)) & (examples/sec) \\  & & (Neg. Log Perp.) & (hours) & \\ \hline T5-Base & — & -1.731 & Not achieved\({}^{\dagger}\) & 1600 \\ T5-Large & — & -1.550 & 131.1 & 470 \\ \hline MoE-Base & 2.0 & -1.547 & 68.7 & 840 \\ Switch-Base & 2.0 & -1.554 & 72.8 & 860 \\ \hline MoE-Base & 1.25 & -1.559 & 80.7 & 790 \\ Switch-Base & 1.25 & -1.553 & 65.0 & 910 \\ \hline MoE-Base & 1.0 & -1.572 & 80.1 & 860 \\ Switch-Base & 1.0 & -1.561 & **62.8** & 1000 \\ Switch-Base+ & 1.0 & **-1.534** & 67.6 & 780 \\ \hline \hline \end{tabular}
+\end{table}
+Table 1: Benchmarking Switch versus MoE. Head-to-head comparison measuring per step and per time benefits of the Switch Transformer over the MoE Transformer and T5 dense baselines. We measure quality by the negative log perplexity and the time to reach an arbitrary chosen quality threshold of Neg. Log Perp.=-1.50. All MoE and Switch Transformer models use 128 experts, with experts at every other feed-forward layer. For Switch-Base+, we increase the model size until it matches the speed of the MoE model by increasing the model hidden-size from 768 to 896 and the number of heads from 14 to 16. All models are trained with the same amount of computation (32 cores) and on the same hardware (TPUv3). Further note that all our models required pre-training beyond 100k steps to achieve our level threshold of -1.50. \(\dagger\) T5-Base did not achieve this negative log perplexity in the 100k steps the models were trained.
+
+are broadcast through all-to-all communication operations, but we still benefit from the increased stability of float32.
+
+**Smalller parameter initialization for stability**. Appropriate initialization is critical to successful training in deep learning and we especially observe this to be true for Switch Transformer. We initialize our weight matrices by drawing elements from a truncated normal distribution with mean \(\mu=0\) and standard deviation \(\sigma=\sqrt{s/n}\) where \(s\) is a scale hyper-parameter and \(n\) is the number of input units in the weight tensor (e.g. fan-in).6
+
+Footnote 6: Values greater than two standard deviations from the mean are resampled.
+
+As an additional remedy to the instability, we recommend reducing the default Transformer initialization scale \(s=1.0\) by a factor of 10. This both improves quality and reduces the likelihood of destabilized training in our experiments. Table 3 measures the improvement of the model quality and reduction of the variance early in training.
+
+We find that the average model quality, as measured by the Neg. Log Perp., is dramatically improved and there is a far reduced variance across runs. Further, this same initialization scheme is broadly effective for models spanning several orders of magnitude. We use the same approach to stably train models as small as our 223M parameter baseline to enormous models in excess of one tr
+
+\begin{table}
+\begin{tabular}{c c c} \hline \hline Model & Quality & Speed \\ (precision) & (Neg. Log Perp.) (\(\uparrow\)) & (Examples/sec) (\(\uparrow\)) \\ \hline Switch-Base (float32) & -1.718 & 1160 \\ Switch-Base (bfloat16) & -3.780 [_diverged_] & **1390** \\ Switch-Base (Selective precision) & **-1.716** & 1390 \\ \hline \hline \end{tabular}
+\end{table}
+Table 2: Selective precision. We cast the local routing operations to float32 while preserving bfloat16 precision elsewhere to stabilize our model while achieving nearly equal speed to (unstable) bfloat16-precision training. We measure the quality of a 32 expert model after a fixed step count early in training its speed performance. For both Switch-Base in float32 and with Selective precision we notice similar learning dynamics.
+
+\begin{table}
+\begin{tabular}{c c c} \hline \hline Model (Initialization scale) & Average Quality & Std. Dev. of Quality \\  & (Neg. Log Perp.) & (Neg. Log Perp.) \\ \hline Switch-Base (0.1x-init) & **-2.72** & **0.01** \\ Switch-Base (1.0x-init) & -3.60 & 0.68 \\ \hline \hline \end{tabular}
+\end{table}
+Table 3: Reduced initialization scale improves stability. Reducing the initialization scale results in better model quality and more stable training of Switch Transformer. Here we record the average and standard deviation of model quality, measured by the negative log perplexity, of a 32 expert model after 3.5k steps (3 random seeds each).
+
+**Regularizing large sparse models.** Our paper considers the common NLP approach of pre-training on a large corpus followed by fine-tuning on smaller downstream tasks such as summarization or question answering. One issue that naturally arises is overfitting since many fine-tuning tasks have very few examples. During fine-tuning of standard Transformers, Raffel et al. (2019) use dropout (Srivastava et al., 2014) at each layer to prevent overfitting. Our Switch Transformers have significantly more parameters than the FLOP matched dense baseline, which can lead to more severe overfitting on these smaller downstream tasks.
+
+We thus propose a simple way to alleviate this issue during fine-tuning: increase the dropout inside the experts, which we name as _expert dropout_. During fine-tuning we simply increase the dropout rate by a significant amount only at the interim feed-forward computation at each expert layer. Table 4 has the results for our expert dropout protocol. We observe that simply increasing the dropout across all layers leads to worse performance. However, setting a smaller dropout rate (0.1) at non-expert layers and a much larger dropout rate (0.4) at expert layers leads to performance improvements on four smaller downstream tasks.
+
+## 3 Scaling Properties
+
+We present a study of the _scaling properties_ of the Switch Transformer architecture during pre-training. Per Kaplan et al. (2020), we consider a regime where the model is not bottlenecked by either the computational budget or amount of data. To avoid the data bottleneck, we use the large C4 corpus with over 180B target tokens (Raffel et al., 2019) and we train until diminishing returns are observed.
+
+The number of experts is the most efficient dimension for scaling our model. Increasing the experts keeps the computational cost approximately fixed since the model only selects one expert per token, regardless of the number of experts to choose from. The router must compute a probability distribution over more experts, however, this is a lightweight computation of cost \(O(d_{model}\times\text{num experts})\) where \(d_{model}\) is the embedding dimension of
+
+\begin{table}
+\begin{tabular}{c c c c c} \hline \hline Model (dropout) & GLUE & CNNDM & SQuAD & SuperGLUE \\ \hline T5-Base (d=0.1) & 82.9 & **19.6** & 83.5 & 72.4 \\ Switch-Base (d=0.1) & 84.7 & 19.1 & **83.7** & **73.0** \\ Switch-Base (d=0.2) & 84.4 & 19.2 & **83.9** & **73.2** \\ Switch-Base (d=0.3) & 83.9 & 19.6 & 83.4 & 70.7 \\ Switch-Base (d=0.1, ed=0.4) & **85.2** & **19.6** & **83.7** & **73.0** \\ \hline \hline \end{tabular}
+\end{table}
+Table 4: Fine-tuning regularization results. A sweep of dropout rates while fine-tuning Switch Transformer models pre-trained on 34B tokens of the C4 data set (higher numbers are better). We observe that using a lower standard dropout rate at all non-expert layer, with a much larger dropout rate on the expert feed-forward layers, to perform the best.
+
+tokens passed between the layers. In this section, we consider the scaling properties on a step-basis and a time-basis with a fixed computational budget.
+
+### Scaling Results on a Step-Basis
+
+Figure 4 demonstrates consistent scaling benefits with the number of experts when training all models for a fixed number of steps. We observe a clear trend: when keeping the FLOPS per token fixed, having more parameters (experts) speeds up training. The left Figure demonstrates consistent scaling properties (with fixed FLOPS per token) between sparse model parameters and test loss. This reveals the advantage of scaling along this additional axis of sparse model parameters. Our right Figure measures sample efficiency of a dense model variant and four FLOP-matched sparse variants. We find that increasing the number of experts leads to more sample efficient models. Our Switch-Base 64 expert model achieves the same performance of the T5-Base model at step 60k at step 450k, which is a 7.5x speedup in terms of step time. In addition, consistent with the findings of Kaplan et al. (2020), we find that larger models are also more _sample efficient_--learning more quickly for a fixed number of observed tokens.
+
+Figure 4: Scaling properties of the Switch Transformer. Left Plot: We measure the quality improvement, as measured by perplexity, as the parameters increase by scaling the number of experts. The top-left point corresponds to the T5-Base model with 223M parameters. Moving from top-left to bottom-right, we double the number of experts from 2, 4, 8 and so on until the bottom-right point of a 256 expert model with 14.7B parameters. Despite all models using an equal computational budget, we observe consistent improvements scaling the number of experts. Right Plot: Negative log perplexity per step sweeping over the number of experts. The dense baseline is shown with the purple line and we note improved sample efficiency of our Switch-Base models.
+
+### Scaling Results on a Time-Basis
+
+Figure 4 demonstrates that on a step basis, as we increase the number of experts, the performance consistently improves. While our models have roughly the same amount of FLOPS per token as the baseline, our Switch Transformers incurs additional communication costs across devices as well as the extra computation of the routing mechanism. Therefore, the increased sample efficiency observed on a step-basis doesn't necessarily translate to a better model quality as measured by wall-clock. This raises the question:
+
+_For a fixed training duration and computational budget, should one train a dense or a sparse model?_
+
+Figures 5 and 6 address this question. Figure 5 measures the pre-training model quality as a function of time. For a fixed training duration and computational budget, Switch Transformers yield a substantial speed-up. In this setting, our Switch-Base 64 expert model trains in _one-seventh_ the time that it would take the T5-Base to get similar perplexity.
+
+### Scaling Versus a Larger Dense Model
+
+The above analysis shows that a computationally-matched dense model is outpaced by its Switch counterpart. Figure 6 considers a different scenario: what if we instead had allocated our resources to a larger dense model? We do so now, measuring Switch-Base against the next strong baseline, _T5-Large_. But despite T5-Large applying 3.5x more FLOPs per token,
+
+Figure 5: Speed advantage of Switch Transformer. All models trained on 32 TPUv3 cores with equal FLOPs per example. For a fixed amount of computation and training time, Switch Transformers significantly outperform the dense Transformer baseline. Our 64 expert Switch-Base model achieves the same quality in _one-seventh_ the time of the T5-Base and continues to improve.
+
+Switch-Base is still more sample efficient and yields a 2.5x speedup. Furthermore, more gains can be had simply by designing a new, larger sparse version, Switch-Large, which is FLOP-matched to T5-Large. We do this and demonstrate superior scaling and fine-tuning in the following section.
+
+## 4 Downstream Results
+
+Section 3 demonstrated the superior scaling properties while pre-training, but we now validate that these gains translate to improved language learning abilities on downstream tasks. We begin by fine-tuning on a diverse set of NLP tasks. Next we study reducing the memory footprint of our sparse models by over 90% by distilling into small--and easily deployed--dense baselines. Finally, we conclude this section measuring the improvements in a multi-task, multilingual setting, where we show that Switch Transformers are strong multi-task learners, improving over the multilingual T5-base model across all 101 languages.
+
+### Fine-Tuning
+
+**Baseline and Switch models used for fine-tuning.** Our baselines are the highly-tuned 223M parameter T5-Base model and the 739M parameter T5-Large model (Raffel et al., 2019). For both versions, we design a FLOP-matched Switch Transformer, with many more parameters, which is summarized in Table 9.7 Our baselines differ slightly from those in Raffel et al. (2019) because we pre-train on an improved C4 corpus which removes intra-example text duplication and thus increases the efficacy as a pre-training task Lee et al.
+
+Figure 6: Scaling Transformer models with Switch layers or with standard dense model scaling. Left Plot: Switch-Base is more sample efficient than both the T5-Base, and T5-Large variant, which applies 3.5x more FLOPS per token. Right Plot: As before, on a wall-clock basis, we find that Switch-Base is still faster, and yields a 2.5x speedup over T5-Large.
+
+(2021). In our protocol we pre-train with \(2^{20}\) (1,048,576) tokens per batch for 550k steps amounting to 576B total tokens. We then fine-tune across a diverse set of tasks using a dropout rate of 0.1 for all layers except the Switch layers, which use a dropout rate of 0.4 (see Table 4). We fine-tune using a batch-size of 1M for 16k steps and for each task, we evaluate model quality every 200-steps and report the peak performance as computed on the validation set.
+
+**Fine-tuning tasks and data sets.** We select tasks probing language capabilities including question answering, summarization and knowledge about the world. The language benchmarks GLUE (Wang et al., 2018) and SuperGLUE (Wang et al., 2019) are handled as composite mixtures with all the tasks blended in proportion to the amount of tokens present in each. These benchmarks consist of tasks requiring sentiment analysis (SST-2), word sense disambiguation (WIC), sentence similarty (MRPC, STS-B, QQP), natural language inference (MNLI, QNLI, RTE, CB), question answering (MultiRC, RECORD, BoolQ), coreference resolution (WNLI, WSC) and sentence completion (COPA) and sentence acceptability (CoLA). The CNNDM (Hermann et al., 2015) and BBC XSum (Narayan et al., 2018) data sets are used to measure the ability to summarize articles. Question answering is probed with the SQuAD data set (Rajpurkar et al., 2016) and the ARC Reasoning Challenge (Clark et al., 2018). And as in Roberts et al. (2020), we evaluate the knowledge of our models by fine-tuning on three closed-book question answering data sets: Natural Questions (Kwiatkowski et al., 2019), Web Questions (Berant et al., 2013) and Trivia QA (Joshi et al., 2017). Closed-book refers to questions posed with no supplemental reference or context material. To gauge the model's common sense reasoning we evaluate it on the Winogrande Schema Challenge (Sakaguchi et al., 2020). And finally, we test our model's natural language inference capabilities on the Adversarial NLI Benchmark (Nie et al., 2019).
+
+**Fine-tuning metrics.** The following evaluation metrics are used throughout the paper: We report the average scores across all subtasks for GLUE and SuperGLUE. The Rouge-2 metric is used both the CNNDM and XSum. In SQuAD and the closed book tasks (Web, Natural, and Trivia Questions) we report the percentage of answers exactly matching the target (refer to Roberts et al. (2020) for further details and deficiency of this measure). Finally, in ARC Easy, ARC Challenge, ANLI, and Winogrande we report the accuracy of the generated responses.
+
+**Fine-tuning results.** We observe significant downstream improvements across many natural language tasks. Notable improvements come from SuperGLUE, where we find FLOP-matched Switch variants improve by 4.4 and 2 percentage points over the T5-Base and T5-Large baselines, respectively as well as large improvements in Winogrande, closed book Trivia QA, and XSum.8 In our fine-tuning study, the only tasks where we do not observe gains are on the AI2 Reasoning Challenge (ARC) data sets where the T5-Base outperforms Switch-Base on the challenge data set and T5-Large outperforms Switch-Large on the easy data set. Taken as a whole, we observe significant improvements spanning both reasoning and knowledge-heavy tasks. This validates our architecture, not just as one that pre-trains well, but can translate quality improvements to downstream tasks via fine-tuning.
+
+### Distillation
+
+Deploying massive neural networks with billions, or trillions, of parameters is inconvenient. To alleviate this, we study distilling (Hinton et al., 2015) large sparse models into small dense models. Future work could additionally study distilling large models into smaller _sparse_ models.
+
+**Distillation techniques.** In Table 6 we study a variety of distillation techniques. These techniques are built off of Sanh et al. (2019), who study distillation methods for BERT models. We find that initializing the dense model with the non-expert weights yields a modest improvement. This is possible since all models are FLOP matched, so non-expert layers will have the same dimensions. Since expert layers are usually only added at every or every other FFN layer in a Transformer, this allows for many of the weights to be initialized with trained parameters. Furthermore, we observe a distillation improvement using a mixture of 0.25 for the teacher probabilities and 0.75 for the ground truth label. By combining both techniques we preserve \(\approx 30\%\) of the quality gains from the larger sparse models with only \(\approx 1/20^{th}\) of the parameters. The quality gain refers to the percent of
+
+\begin{table}
+\begin{tabular}{c c c c c} \hline \hline Model & GLUE & SQuAD & SuperGLUE & Winogrande (XL) \\ \hline T5-Base & 84.3 & 85.5 & 75.1 & 66.6 \\ Switch-Base & **86.7** & **87.2** & **79.5** & **73.3** \\ T5-Large & 87.8 & 88.1 & 82.7 & 79.1 \\ Switch-Large & **88.5** & **88.6** & **84.7** & **83.0** \\ \hline \hline Model & XSum & ANLI (R3) & ARC Easy & ARC Chal. \\ \hline T5-Base & 18.7 & 51.8 & 56.7 & **35.5** \\ Switch-Base & **20.3** & **54.0** & **61.3** & 32.8 \\ T5-Large & 20.9 & 56.6 & **68.8** & **35.5** \\ Switch-Large & **22.3** & **58.6** & 66.0 & **35.5** \\ \hline \hline Model & CB Web QA & CB Natural QA & CB Trivia QA & \\ \hline T5-Base & 26.6 & 25.8 & 24.5 & \\ Switch-Base & **27.4** & **26.8** & **30.7** & \\ T5-Large & 27.7 & 27.6 & 29.5 & \\ Switch-Large & **31.3** & **29.5** & **36.9** & \\ \hline \hline \end{tabular}
+\end{table}
+Table 5: Fine-tuning results. Fine-tuning results of T5 baselines and Switch models across a diverse set of natural language tests (validation sets; higher numbers are better). We compare FLOP-matched Switch models to the T5-Base and T5-Large baselines. For most tasks considered, we find significant improvements of the Switch-variants. We observe gains across both model sizes and across both reasoning and knowledge-heavy language tasks.
+
+the quality difference between Switch-Base (Teacher) and T5-Base (Student). Therefore, a quality gain of 100% implies the Student equals the performance of the Teacher.
+
+**Achievable compression rates.** Using our best distillation technique described in Table 6, we distill a wide variety of sparse models into dense models. We distill Switch-Base versions, sweeping over an increasing number of experts, which corresponds to varying between 1.1B to 14.7B parameters. Through distillation, we can preserve 37% of the quality gain of the 1.1B parameter model while compressing 82%. At the extreme, where we compress the model 99%, we are still able to maintain 28% of the teacher's model quality improvement.
+
+**Distilling a fine-tuned model.** We conclude this with a study of distilling a fine-tuned sparse model into a dense model. Table 8 shows results of distilling a 7.4B parameter Switch-Base model, fine-tuned on the SuperGLUE task, into the 223M T5-Base. Similar to our pre-training results, we find we are able to preserve 30% of the gains of the sparse model when distilling into a FLOP matched dense variant. One potential future avenue, not considered here, may examine the specific experts being used for fine-tuning tasks and extracting them to achieve better model compression.
+
+### Multilingual Learning
+
+In our final set of downstream experiments, we measure the model quality and speed trade-offs while pre-training on a mixture of 101 different languages. We build and benchmark off the recent work of mT5 (Xue et al., 2020), a multilingual extension to T5. We pre-train on the multilingual variant of the Common Crawl data set (mC4) spanning 101 languages introduced in mT5, but due to script variants within certain languages, the mixture contains 107 tasks.
+
+In Figure 7 we plot the quality improvement in negative log perplexity for all languages of a FLOP-matched Switch model, mSwitch-Base to the T5 base variant, mT5-Base. After
+
+\begin{table}
+\begin{tabular}{l r r} \hline \hline Technique & Parameters & Quality (\(\uparrow\)) \\ \hline T5-Base & 223M & -1.636 \\ Switch-Base & 3,800M & -1.444 \\ \hline Distillation & 223M & (3\%) -1.631 \\ + Init. non-expert weights from teacher & 223M & (20\%) -1.598 \\ + 0.75 mix of hard and soft loss & 223M & (29\%) -1.580 \\ \hline Initialization Baseline (no distillation) & & \\ Init. non-expert weights from teacher & 223M & -1.639 \\ \hline \hline \end{tabular}
+\end{table}
+Table 6: Distilling Switch Transformers for Language Modeling. Initializing T5-Base with the non-expert weights from Switch-Base and using a loss from a mixture of teacher and ground-truth labels obtains the best performance. We can distill 30% of the performance improvement of a large sparse model with 100x more parameters back into a small dense model. For a final baseline, we find no improvement of T5-Base initialized with the expert weights, but trained normally without distillation.
+
+pre-training both versions for 1M steps, we find that on _all_ 101 languages considered, Switch Transformer increases the final negative log perplexity over the baseline. In Figure 8, we present a different view and now histogram the per step _speed-up_ of using Switch Transformer over the mT5-Base.9 We find a mean speed-up over mT5-Base of 5x and that 91% of languages achieve at least a 4x speedup. This presents evidence that Switch Transformers are effective multi-task and multi-lingual learners.
+
+Footnote 9: The speedup on a step basis is computed as the ratio of the number of steps for the baseline divided by the number of steps required by our model to reach that same quality.
+
+## 5 Designing Models with Data, Model, and Expert-Parallelism
+
+Arbitrarily increasing the number of experts is subject to diminishing returns (Figure 4). Here we describe _complementary_ scaling strategies. The common way to scale a Transformer is to increase dimensions in tandem, like \(d_{model}\) or \(d_{ff}\). This increases both the parameters
+
+\begin{table}
+\begin{tabular}{c c|c c c c c} \hline \hline  & Dense & \multicolumn{5}{c}{Sparse} \\ \hline Parameters & 223M & 1.1B & 2.0B & 3.8B & 7.4B & 14.7B \\ \hline Pre-trained Neg. Log Perp. (\(\uparrow\)) & -1.636 & -1.505 & -1.474 & -1.444 & -1.432 & -1.427 \\ Distilled Neg. Log Perp. (\(\uparrow\)) & — & -1.587 & -1.585 & -1.579 & -1.582 & -1.578 \\ Percent of Teacher Performance & — & 37\% & 32\% & 30 \% & 27 \% & 28 \% \\ Compression Percent & — & 82 \% & 90 \% & 95 \% & 97 \% & 99 \% \\ \hline \hline \end{tabular}
+\end{table}
+Table 7: Distillation compression rates. We measure the quality when distilling large sparse models into a dense baseline. Our baseline, T5-Base, has a -1.636 Neg. Log Perp. quality. In the right columns, we then distill increasingly large sparse models into this same architecture. Through a combination of weight-initialization and a mixture of hard and soft losses, we can shrink our sparse teachers by 95%+ while preserving 30% of the quality gain. However, for significantly better and larger pre-trained teachers, we expect larger student models would be necessary to achieve these compression rates.
+
+\begin{table}
+\begin{tabular}{c c c|c} \hline \hline Model & Parameters & FLOPS & SuperGLUE (\(\uparrow\)) \\ \hline T5-Base & 223M & 124B & 74.6 \\ Switch-Base & 7410M & 124B & 81.3 \\ Distilled T5-Base & 223M & 124B & (30\%) 76.6 \\ \hline \hline \end{tabular}
+\end{table}
+Table 8: Distilling a fine-tuned SuperGLUE model. We distill a Switch-Base model fine-tuned on the SuperGLUE tasks into a T5-Base model. We observe that on smaller data sets our large sparse model can be an effective teacher for distillation. We find that we again achieve 30% of the teacher’s performance on a 97% compressed model.
+
+and computation performed and is ultimately limited by the memory per accelerator. Once it exceeds the size of the accelerator's memory, single program multiple data (SPMD) model-parallelism can be employed. This section studies the trade-offs of combining data, model, and expert-parallelism.
+
+**Reviewing the Feed-Forward Network (FFN) Layer.** We use the FFN layer as an example of how data, model and expert-parallelism works in Mesh TensorFlow (Shazeer et al., 2018) and review it briefly here. We assume \(B\) tokens in the batch, each of dimension
+
+Figure 8: Multilingual pre-training on 101 languages. We histogram for each language, the step speedup of Switch Transformers over the FLOP matched T5 dense baseline to reach the same quality. Over all 101 languages, we achieve a mean step speed-up over mT5-Base of 5x and, for 91% of languages, we record a 4x, or greater, speedup to reach the final perplexity of mT5-Base.
+
+Figure 7: Multilingual pre-training on 101 languages. Improvements of Switch T5 Base model over dense baseline when multi-task training on 101 languages. We observe Switch Transformers to do quite well in the multi-task training setup and yield improvements on all 101 languages.
+
+\(d_{model}\). Both the input (\(x\)) and output (\(y\)) of the FFN are of size [\(B\), \(d_{model}\)] and the intermediate (\(h\)) is of size [\(B\), \(d_{ff}\)] where \(d_{ff}\) is typically several times larger than \(d_{model}\). In the FFN, the intermediate is \(h=xW_{in}\) and then the output of the layer is \(y=ReLU(h)W_{out}\). Thus \(W_{in}\) and \(W_{out}\) are applied independently to each token and have sizes [\(d_{model}\), \(d_{ff}\)] and [\(d_{ff}\), \(d_{model}\)].
+
+We describe two aspects of partitioning: how the _weights_ and _batches of data_ divide over cores, depicted in Figure 9. We denote all cores available as \(N\) which Mesh Tensorflow may then remap into a logical multidimensional mesh of processors. Here we create a two-dimensional logical mesh, with one dimension representing the number of ways for data-parallel sharding (\(n\)) and the other, the model-parallel sharding (\(m\)). The total cores must equal the ways to shard across both data and model-parallelism, e.g. \(N=n\times m\). To shard the layer across cores, the tensors containing that batch of \(B\) tokens are sharded across \(n\) data-parallel cores, so each core contains \(B/n\) tokens. Tensors and variables with \(d_{ff}\) are then sharded across \(m\) model-parallel cores. For the variants with experts-layers, we consider \(E\) experts, each of which can process up to \(C\) tokens.
+
+### Data Parallelism
+
+When training data parallel models, which is the standard for distributed training, then all cores are allocated to the data-parallel dimension or \(n=N,m=1\). This has the advantage that no communication is needed until the entire forward and backward pass is finished and the gradients need to be then aggregated across all cores. This corresponds to the left-most column of Figure 9.
+
+### Model Parallelism
+
+We now consider a scenario where all cores are allocated exclusively to the model-parallel dimension and so \(n=1,m=N\). Now all cores must keep the full \(B\) tokens and each core will contain a unique slice of the weights. For each forward and backward pass, a communication cost is now incurred. Each core sends a tensor of [\(B\), \(d_{model}\)] to compute the second matrix multiplication \(ReLU(h)W_{out}\) because the \(d_{ff}\) dimension is partitioned and must be summed over. As a general rule, whenever a dimension that is partitioned across cores must be summed, then an all-reduce operation is added for both the forward and backward pass. This contrasts with pure data parallelism where an all-reduce only occurs at the end of the entire forward and backward pass.
+
+### Model and Data Parallelism
+
+It is common to mix both model and data parallelism for large scale models, which was done in the largest T5 models (Raffel et al., 2019; Xue et al., 2020) and in GPT-3 (Brown et al., 2020). With a total of \(N=n\times m\) cores, now each core will be responsible for \(B/n\) tokens and \(d_{ff}/m\) of both the weights and intermediate activation. In the forward and backward pass each core communicates a tensor of size \([B/n,d_{model}]\) in an all-reduce operation.
+
+Figure 9: Data and weight partitioning strategies. Each 4\(\times\)4 dotted-line grid represents 16 cores and the shaded squares are the data contained on that core (either model weights or batch of tokens). We illustrate both how the model weights and the data tensors are split for each strategy. **First Row:** illustration of how _model weights_ are split across the cores. Shapes of different sizes in this row represent larger weight matrices in the Feed Forward Network (FFN) layers (e.g larger \(d_{ff}\) sizes). Each color of the shaded squares identifies a unique weight matrix. The number of parameters _per core_ is fixed, but larger weight matrices will apply more computation to each token. **Second Row:** illustration of how the _data batch_ is split across cores. Each core holds the same number of tokens which maintains a fixed memory usage across all strategies. The partitioning strategies have different properties of allowing each core to either have the same tokens or different tokens across cores, which is what the different colors symbolize.
+
+### Expert and Data Parallelism
+
+Next we describe the partitioning strategy for expert and data parallelism. Switch Transformers will allocate all of their cores to the data partitioning dimension \(n\), which will also correspond to the number of experts in the model. For each token per core a router locally computes assignments to the experts. The output is a binary matrix of size [\(n\), \(B/n\), \(E\), \(C\)] which is partitioned across the first dimension and determines expert assignment. This binary matrix is then used to do a gather via matrix multiplication with the input tensor of [\(n\), \(B/n\), \(d_{model}\)].
+
+\[\text{einsum}([n,B/n,d_{model}],[n,B/n,E,C],\text{dimension}=[B/n]) \tag{7}\]
+
+resulting in the final tensor of shape [\(n\), \(E\), \(C\), \(d_{model}\)], which is sharded across the first dimension. Because each core has its own expert, we do an all-to-all communication of size [\(E\), \(C\), \(d_{model}\)] to now shard the \(E\) dimension instead of the \(n\)-dimension. There are additional communication costs of bfloat16 tensors of size \(E\times C\times d_{model}\) in the forward pass to analogously receive the tokens from each expert located on different cores. See Appendix F for a detailed analysis of the expert partitioning code.
+
+### Expert, Model and Data Parallelism
+
+In the design of our best model, we seek to balance the FLOPS per token and the parameter count. When we scale the number of experts, we increase the number of parameters, but do not change the FLOPs per token. In order to increase FLOPs, we must also increase the \(d_{ff}\) dimension (which also increases parameters, but at a slower rate). This presents a trade-off: as we increase \(d_{ff}\) we will run out of memory per core, which then necessitates increasing \(m\). But since we have a fixed number of cores \(N\), and \(N=n\times m\), we must decrease \(n\), which forces use of a smaller batch-size (in order to hold tokens per core constant).
+
+When combining both model and expert-parallelism, we will have all-to-all communication costs from routing the tokens to the correct experts along with the internal all-reduce communications from the model parallelism. Balancing the FLOPS, communication costs and memory per core becomes quite complex when combining all three methods where the best mapping is empirically determined. See our further analysis in section 5.6 for how the number of experts effects the downstream performance as well.
+
+### Towards Trillion Parameter Models
+
+Combining expert, model and data parallelism, we design two large Switch Transformer models, one with 395 billion and 1.6 trillion parameters, respectively. We study how these models perform on both up-stream pre-training as language models and their downstream fine-tuning performance. The parameters, FLOPs per sequence and hyper-parameters of the two different models are listed below in Table 9. Standard hyper-parameters of the Transformer, including \(d_{model}\), \(d_{ff}\), \(d_{kv}\), number of heads and number of layers are described, as well as a less common feature, \(FFN_{GEGLU}\), which refers to a variation of the FFN layer where the expansion matrix is substituted with two sets of weights which are non-linearly combined (Shazeer, 2020).
+
+The Switch-C model is designed using only expert-parallelism, and no model-parallelism, as described earlier in Section 5.4. As a result, the hyper-parameters controlling the width,depth, number of heads, and so on, are all much smaller than the T5-XXL model. In contrast, the Switch-XXL is FLOP-matched to the T5-XXL model, which allows for larger dimensions of the hyper-parameters, but at the expense of additional communication costs induced by model-parallelism (see Section 5.5 for more details).
+
+**Sample efficiency versus T5-XXL.** In the final two columns of Table 9 we record the negative log perplexity on the C4 corpus after 250k and 500k steps, respectively. After 250k steps, we find both Switch Transformer variants to improve over the T5-XXL version's negative log perplexity by over 0.061.10 To contextualize the significance of a gap of 0.061, we note that the T5-XXL model had to train for an _additional_ 250k steps to increase 0.052. The gap continues to increase with additional training, with the Switch-XXL model out-performing the T5-XXL by 0.087 by 500k steps.
+
+Footnote 10: This reported quality difference is a lower bound, and may actually be larger. The T5-XXL was pre-trained on an easier C4 data set which included duplicated, and thus easily copied, snippets within examples.
+
+**Training instability.** However, as described in the introduction, large sparse models can be unstable, and as we increase the scale, we encounter some sporadic issues. We find that the larger Switch-C model, with 1.6T parameters and 2048 experts, exhibits no training instability at all. Instead, the Switch XXL version, with nearly 10x larger FLOPs per sequence, is sometimes unstable. As a result, though this is our better model on a step-basis, we do not pre-train for a full 1M steps, in-line with the final reported results of T5 (Raffel et al., 2019).
+
+\begin{table}
+\begin{tabular}{c|c c c c c c c} \hline \hline Model & Parameters & FLOPs/seq & \(d_{\text{mubd}}\) & \(FFN_{\text{CEGLU}}\) & \(d_{ff}\) & \(d_{\text{ks}}\) & Num. Heads \\ \hline T5-Base & 0.2B & 124B & 768 & ✓ & 2048 & 64 & 12 \\ T5-Large & 0.7B & 425B & 1024 & ✓ & 2816 & 64 & 16 \\ T5-XXL & 11B & 6.3T & 4096 & ✓ & 10240 & 64 & 64 \\ \hline Switch-Base & 7B & 124B & 768 & ✓ & 2048 & 64 & 12 \\ Switch-Large & 26B & 425B & 1024 & ✓ & 2816 & 64 & 16 \\ Switch-XXL & 395B & 6.3T & 4096 & ✓ & 10240 & 64 & 64 \\ Switch-C & 1571B & 890B & 2080 & & 6144 & 64 & 32 \\ \hline \hline Model & Expert Freq. & Num. Layers & Num Experts & Neg. Log Perp. @250k & Neg. Log Perp. @ 500k & \\ \hline T5-Base & – & 12 & – & -1.599 & -1.556 & \\ T5-Large & – & 24 & – & -1.402 & -1.350 & \\ T5-XXL & – & 24 & – & -1.147 & -1.095 & \\ \hline Switch-Base & 1/2 & 12 & 128 & -1.370 & -1.306 & \\ Switch-Large & 1/2 & 24 & 128 & -1.248 & -1.177 & \\ Switch-XXL & 1/2 & 24 & 64 & -**1.086** & **-1.008** & \\ Switch-C & 1 & 15 & 2048 & -1.096 & -1.043 & \\ \hline \hline \end{tabular}
+\end{table}
+Table 9: Switch model design and pre-training performance. We compare the hyper-parameters and pre-training performance of the T5 models to our Switch Transformer variants. The last two columns record the pre-training model quality on the C4 data set after 250k and 500k steps, respectively. We observe that the Switch-C Transformer variant is 4x faster to a fixed perplexity (with the same compute budget) than the T5-XXL model, with the gap increasing as training progresses.
+
+**Reasoning fine-tuning performance.** As a preliminary assessment of the model quality, we use a Switch-XXL model partially pre-trained on 503B tokens, or approximately half the text used by the T5-XXL model. Using this checkpoint, we conduct multi-task training for efficiency, where all tasks are learned jointly, rather than individually fine-tuned. We find that SQuAD accuracy on the validation set increases to 89.7 versus state-of-the-art of 91.3. Next, the average SuperGLUE test score is recorded at 87.5 versus the T5 version obtaining a score of 89.3 compared to the state-of-the-art of 90.0 (Wang et al., 2019). On ANLI (Nie et al., 2019), Switch XXL improves over the prior state-of-the-art to get a 65.7 accuracy versus the prior best of 49.4 (Yang et al., 2020). We note that while the Switch-XXL has state-of-the-art Neg. Log Perp. on the upstream pre-training task, its gains have not yet fully translated to SOTA downstream performance. We study this issue more in Appendix E.
+
+**Knowledge-based fine-tuning performance.** Finally, we also conduct an early examination of the model's knowledge with three closed-book knowledge-based tasks: Natural Questions, WebQuestions and TriviaQA, without additional pre-training using Salient Span Masking (Guu et al., 2020). In all three cases, we observe improvements over the prior state-of-the-art T5-XXL model (without SSM). Natural Questions exact match increases to 34.4 versus the prior best of 32.8, Web Questions increases to 41.0 over 37.2, and TriviaQA increases to 47.5 versus 42.9.
+
+Summing up, despite training on less than half the data of other models, we already find comparable, and sometimes state-of-the-art, model quality. Currently, the Switch Transformer translates substantial upstream gains better to knowledge-based tasks, than reasoning-tasks (see Appendix E). Extracting stronger fine-tuning performance from large expert models is an active research question, and the pre-training perplexity indicates future improvements should be possible.
+
+## 6 Related Work
+
+The importance of scale in neural networks is widely recognized and several approaches have been proposed. Recent works have scaled models to billions of parameters through using model parallelism (e.g. splitting weights and tensors across multiple cores) (Shazeer et al., 2018; Rajbhandari et al., 2019; Raffel et al., 2019; Brown et al., 2020; Shoeybi et al., 2019). Alternatively, Harlap et al. (2018); Huang et al. (2019) propose using pipeline based model parallelism, where different layers are split across devices and micro-batches are _pipelined_ to the different layers. Finally, Product Key networks (Lample et al., 2019) were proposed to scale up the capacity of neural networks by doing a lookup for learnable embeddings based on the incoming token representations to a given layer.
+
+Our work studies a specific model in a class of methods that do _conditional_ computation, where computation decisions are made dynamically based on the input. Cho and Bengio (2014) proposed adaptively selecting weights based on certain bit patterns occuring in the model hidden-states. Eigen et al. (2013) built stacked expert layers with dense matrix multiplications and ReLU activations and showed promising results on jittered MNIST and monotone speech. In computer vision Puigcerver et al. (2020) manually route tokens based on semantic classes during upstream pre-training and then select the relevant experts to be used according to the downstream task.
+
+Mixture of Experts (MoE), in the context of modern deep learning architectures, was proven effective in Shazeer et al. (2017). That work added an MoE layer which was stacked between LSTM (Hochreiter and Schmidhuber, 1997) layers, and tokens were separately routed to combinations of experts. This resulted in state-of-the-art results in language modeling and machine translation benchmarks. The MoE layer was reintroduced into the Transformer architecture by the Mesh Tensorflow library (Shazeer et al., 2018) where MoE layers were introduced as a substitute of the FFN layers, however, there were no accompanying NLP results. More recently, through advances in machine learning infrastructure, GShard (Lepikhin et al., 2020), which extended the XLA compiler, used the MoE Transformer to dramatically improve machine translation across 100 languages. Finally Fan et al. (2021) chooses a different deterministic MoE strategy to split the model parameters into non-overlapping groups of languages.
+
+Sparsity along the sequence length dimension (\(L\)) in the Transformer _attention patterns_ has been a successful technique to reduce the attention complexity from \(O(L^{2})\)(Child et al., 2019; Correia et al., 2019; Sukhbaatar et al., 2019; Kitaev et al., 2020; Zaheer et al., 2020; Beltagy et al., 2020). This has enabled learning longer sequences than previously possible. This version of the Switch Transformer does not employ attention sparsity, but these techniques are complimentary, and, as future work, these could be combined to potentially improve learning on tasks requiring long contexts.
+
+## 7 Discussion
+
+We pose and discuss questions about the Switch Transformer, and sparse expert models generally, where sparsity refers to weights, not on attention patterns.
+
+**Isn't Switch Transformer better due to sheer parameter count?** Yes, and by design! Parameters, independent of the total FLOPs used, are a useful axis to scale neural language models. Large models have been exhaustively shown to perform better (Kaplan et al., 2020). But in this case, our model is more sample efficient and faster while using the same computational resources.
+
+**I don't have access to a supercomputer--is this still useful for me?** Though this work has focused on extremely large models, we also find that models with as few as two experts improves performance while easily fitting within memory constraints of commonly available GPUs or TPUs (details in Appendix D). We therefore believe our techniques are useful in small-scale settings.
+
+**Do sparse models outperform dense models on the speed-accuracy Pareto curve?** Yes. Across a wide variety of different models sizes, sparse models outperform dense models per step and on wall clock time. Our controlled experiments show for a fixed amount of computation and time, sparse models outperform dense models.
+
+**I can't deploy a trillion parameter model--can we shrink these models?** We cannot fully preserve the model quality, but compression rates of 10 to 100x are achievable by distilling our sparse models into dense models while achieving \(\approx\)30% of the quality gain of the expert model.
+
+**Why use Switch Transformer instead of a model-parallel dense model?** On a time basis, Switch Transformers can be far more efficient than dense-models with sharded parameters (Figure 6). Also, we point out that this decision is _not_ mutually exclusive--wecan, and do, use model-parallelism in Switch Transformers, increasing the FLOPs per token, but incurring the slowdown of conventional model-parallelism.
+
+**Why aren't sparse models widely used already?** The motivation to try sparse models has been stymied by the massive success of scaling dense models (the success of which is partially driven by co-adaptation with deep learning hardware as argued in Hooker (2020)). Further, sparse models have been subject to multiple issues including (1) model complexity, (2) training difficulties, and (3) communication costs. Switch Transformer makes strides to alleviate these issues.
+
+## 8 Future Work
+
+This paper lays out a simplified architecture, improved training procedures, and a study of how sparse models scale. However, there remain many open future directions which we briefly describe here:
+
+1. A significant challenge is further improving training stability for the largest models. While our stability techniques were effective for our Switch-Base, Switch-Large and Switch-C models (no observed instability), they were not sufficient for Switch-XXL. We have taken early steps towards stabilizing these models, which we think may be generally useful for large models, including using regularizers for improving stability and adapted forms of gradient clipping, but this remains unsolved.
+2. Generally we find that improved pre-training quality leads to better downstream results (Appendix E), though we sometimes encounter striking anomalies. For instance, despite similar perplexities modeling the C4 data set, the 1.6T parameter Switch-C achieves only an 87.7 exact match score in SQuAD, which compares unfavorably to 89.6 for the smaller Switch-XXL model. One notable difference is that the Switch-XXL model applies \(\approx\)10x the FLOPS per token than the Switch-C model, even though it has \(\approx\)4x less unique parameters (395B vs 1.6T). This suggests a poorly understood dependence between fine-tuning quality, _FLOPS per token_ and _number of parameters_.
+3. Perform a comprehensive study of scaling relationships to guide the design of architectures blending data, model and expert-parallelism. Ideally, given the specs of a hardware configuration (computation, memory, communication) one could more rapidly design an optimal model. And, vice versa, this may also help in the design of future hardware.
+4. Our work falls within the family of adaptive computation algorithms. Our approach always used identical, homogeneous experts, but future designs (facilitated by more flexible infrastructure) could support _heterogeneous_ experts. This would enable more flexible adaptation by routing to larger experts when more computation is desired--perhaps for harder examples.
+5. Investigating expert layers outside the FFN layer of the Transformer. We find preliminary evidence that this similarly can improve model quality. In Appendix A, we report quality improvement adding these inside Self-Attention layers, where our layer replaces the weight matrices which produce Q, K, V. However, due to training instabilities with the bfloat16 format, we instead leave this as an area for future work.
+6. Examining Switch Transformer in new and across different modalities. We have thus far only considered language, but we believe that model sparsity can similarly provide advantages in new modalities, as well as multi-modal networks.
+
+This list could easily be extended, but we hope this gives a flavor for the types of challenges that we are thinking about and what we suspect are promising future directions.
+
+## 9 Conclusion
+
+Switch Transformers are scalable and effective natural language learners. We simplify Mixture of Experts to produce an architecture that is easy to understand, stable to train and vastly more sample efficient than equivalently-sized dense models. We find that these models excel across a diverse set of natural language tasks and in different training regimes, including pre-training, fine-tuning and multi-task training. These advances make it possible to train models with hundreds of billion to trillion parameters and which achieve substantial speedups relative to dense T5 baselines. We hope our work motivates sparse models as an effective architecture and that this encourages researchers and practitioners to consider these flexible models in natural language tasks, and beyond.
+
+The authors would like to thank Margaret Li who provided months of key insights into algorithmic improvements and suggestions for empirical studies. Hugo Larochelle for sage advising and clarifying comments on the draft, Irwan Bello for detailed comments and careful revisions, Colin Raffel and Adam Roberts for timely advice on neural language models and the T5 code-base, Yoshua Bengio for advising and encouragement on research in adaptive computation, Jascha Sohl-dickstein for interesting new directions for stabilizing new large scale models and paper revisions, and the Google Brain Team for useful discussions on the paper. Blake Hechtman who provided invaluable help in profiling and improving the training performance of our models.
+
+## Appendix A Switch for Attention
+
+Shazeer et al. (2018); Lepikhin et al. (2020) designed MoE Transformers (Shazeer et al., 2017) by adding MoE layers into the dense feedfoward network (FFN) computations of the Transformer. Similarly, our work also replaced the FFN layer in the Transformer, but we briefly explore here an alternate design. We add Switch layers into the Transformer _Self-Attention_ layers. To do so, we replace the trainable weight matrices that produce the queries, keys and values with Switch layers as seen in Figure 10.
+
+Table 10 records the quality after a fixed number of steps as well as training time for several variants. Though we find improvements, we also found these layers to be more unstable when using bfloat16 precision and thus we did not include them in the final variant.
+
+However, when these layers do train stably, we believe the preliminary positive results suggests a future promising direction.
+
+\begin{table}
+\begin{tabular}{c|c c c c} \hline \hline Model & Precision & Quality & Quality & Speed \\  & & @100k Steps (\(\uparrow\)) & @16H (\(\uparrow\)) & (ex/sec) (\(\uparrow\)) \\ \hline Experts FF & float32 & -1.548 & -1.614 & 1480 \\ Expert Attention & float32 & -1.524 & **-1.606** & 1330 \\ Expert Attention & bfloat16 & [diverges] & [diverges] & – \\ Experts FF + Attention & float32 & **-1.513** & -1.607 & 1240 \\ Expert FF + Attention & bfloat16 & [diverges] & [diverges] & – \\ \hline \hline \end{tabular}
+\end{table}
+Table 10: Switch attention layer results. All models have 32 experts and train with 524k tokens per batch. Experts FF is when experts replace the FFN in the Transformer, which is our standard setup throughout the paper. Experts FF + Attention is when experts are used to replace both the FFN and the Self-Attention layers. When training with bfloat16 precision the models that have experts attention diverge.
+
+Figure 10: Switch layers in attention. We diagram how to incorporate the Switch layer into the Self-Attention transformer block. For each token (here we show two tokens, \(x_{1}\) = “More” and \(x_{2}\) = “Parameters”), one set of weights produces the query and the other set of unique weights produces the shared keys and values. We experimented with each expert being a linear operation, as well as a FFN, as was the case throughout this work. While we found quality improvements using this, we found this to be more unstable when used with low precision number formats, and thus leave it for future work.
+
+## Appendix B Preventing Token Dropping with _No-Token-Left-Behind_
+
+Due to software constraints on TPU accelerators, the shapes of our Tensors must be statically sized. As a result, each expert has a finite and fixed capacity to process token representations. This, however, presents an issue for our model which dynamically routes tokens at run-time that may result in an uneven distribution over experts. If the number of tokens sent to an expert is less than the expert capacity, then the computation may simply be padded - an inefficient use of the hardware, but mathematically correct. However, when the number of tokens sent to an expert is larger than its capacity (expert overflow), a protocol is needed to handle this. Lepikhin et al. (2020) adapts a Mixture-of-Expert model and addresses expert overflow by passing its representation to the next layer without processing through a residual connection which we also follow.
+
+We suspected that having no computation applied to tokens could be very wasteful, especially since if there is overflow on one expert, that means another expert will have extra capacity. With this intuition we create _No-Token-Left-Behind_, which iteratively reroutes any tokens that are at first routed to an expert that is overflowing. Figure 11 shows a graphical description of this method, which will allow us to guarantee almost no tokens will be dropped during training and inference. We hypothesised that this could improve performance and further stabilize training, but we found no empirical benefits. We suspect that once the network learns associations between different tokens and experts, if this association is changed (e.g. sending a token to its second highest expert) then performance could be degraded.
+
+## Appendix C Encouraging Exploration Across Experts
+
+At each expert-layer, the router determines to which expert to send the token. This is a discrete decision over the available experts, conditioned on information about the token's representation. Based on the incoming token representation, the router determines the best expert, however, it receives no counterfactual information about how well it would have done selecting an alternate expert. As in reinforcement learning, a classic exploration-exploitation dilemma arises (Sutton and Barto, 2018). These issues have been similarly noted and addressed differently by Rosenbaum et al. (2017) which demonstrated success in multi-task learning. This particular setting most closely matches that of a contextual bandit (Robbins, 1952). Deterministically selecting the top expert always amounts to an exploitative strategy - we consider balancing exploration to seek better expert assignment.
+
+To introduce exploration, we consider several approaches: 1) deterministic or argmax 2) sampling from the softmax distribution 3) input dropout on the incoming representation 4) multiplicative jitter noise on the incoming representation. The resulting impact on model quality is reported in Table 11. Throughout this work, we use input jitter to inject noise as we have found it to empirically perform the best.
+
+## Appendix D Switch Transformers in Lower Compute Regimes
+
+Switch Transformer is also an effective architecture at small scales as well as in regimes with thousands of cores and trillions of parameters. Many of our prior experiments wereat the scale of 10B+ parameter models, but we show in Figure 12 as few as 2 experts produce compelling gains over a FLOP-matched counterpart. Even if a super computer is not readily available, training Switch Transformers with 2, 4, or 8 experts (as we typically recommend one expert per core) results in solid improvements over T5 dense baselines.
+
+Figure 11: Diagram of the _No-Token-Left-Behind Routing_. Stage 1 is equivalent to Switch routing where tokens are routed to the expert with the highest probability from the router. In Stage 2 we look at all tokens that have overflowed and route them to the expert with which has the second highest probability. Tokens can still be overflowed if their second highest expert has too many tokens, but this allows most of the tokens to be routed. This process can be iterated to guarantee virtually no tokens are dropped at all.
+
+\begin{table}
+\begin{tabular}{c c} \hline Model & Quality (Neg. Log Perp.) (\(\uparrow\)) \\ \hline Argmax & -1.471 \\ Sample softmax & -1.570 \\ Input dropout & -1.480 \\ Input jitter & **-1.468** \\ \hline \end{tabular}
+\end{table}
+Table 11: Router Exploration Strategies. Quality of the Switch Transformer, measured by the negative log perplexity, under different randomness-strategies for selecting the expert (lower is better). There is no material speed performance difference between the variants.
+
+Figure 12: Switch Transformer with few experts. Switch Transformer improves over the baseline even with very few experts. Here we show scaling properties at very small scales, where we improve over the T5-Base model using 2, 4, and 8 experts.
+
+## Appendix E Relation of Upstream to Downstream Model Performance
+
+There is no guarantee that a model's quality on a pre-training objective will translate to downstream task results. Figure 13 presents the correlation of the upstream model quality, for both dense and Switch models, on the C4 pre-training task with two downstream task measures: average SuperGLUE performance and TriviaQA score. We choose these two tasks as one probes the model's reasoning and the other factual knowledge.
+
+We find a consistent correlation, indicating that for both baseline and Switch models, improved pre-training leads to better downstream results. Additionally, for a fixed upstream perplexity we find that both Switch and dense models perform similarly in the small to medium model size regime. However, in the largest model regime (T5-11B/T5-XXL) our largest Switch models, as mentioned in Section 5.6, do not always translate their upstream perplexity well to downstream fine-tuning on the SuperGLUE task. This warrants future investigation and study to fully realize the potential of sparse models. Understanding the fine-tuning dynamics with expert-models is very complicated and is dependent on regularization, load-balancing, and fine-tuning hyper-parameters.
+
+Figure 13: Upstream pre-trained quality to downstream model quality. We correlate the upstream performance with downstream quality on both SuperGLUE and TriviaQA (SOTA recorded without SSM), reasoning and knowledge-heavy benchmarks, respectively (validation sets). We find that, as with the baseline, the Switch model scales with improvements in the upstream pre-training task. For SuperGLUE, we find a loosely linear relation between negative log perplexity and the average SuperGLUE score. However, the dense model often performs better for a fixed perplexity, particularly in the large-scale regime. Conversely, on the knowledge-heavy task, TriviaQA, we find that the Switch Transformer may follow an improved scaling relationship – for a given upstream perplexity, it does better than a dense counterpart. Further statistics (expensive to collect and left to future work) would be necessary to confirm these observations.
+
+## Appendix F Pseudo Code for Switch Transformers
+
+Pseudocode for Switch Transformers in Mesh Tensorflow (Shazeer et al., 2018). No model parallelism is being used for the below code (see 5.4 for more details).
+
+Figure 14: Pseudo code for the load balance loss for Switch Transformers in Mesh Tensorflow.
+
+importmesh_tensorflowasntf defrouter(inputs,capacity_factor):  """Producethecombineanddispatchtensorsusedforsendingand  receivingtokensfromtheirhighestprobabilityexpert."""  #Corelayoutissplitacrossnum_coresforalltensorsandoperations.  #inputsshape:[num_cores,tokens_per_core,d_model] router_weights=ntf.Variable(shape=[d_model,num_experts])
+#router_logitsshape:[num_cores,tokens_per_core,num_experts] router_logits=ntf.einsum([inputs,router_weights],reduced_dim=d_model) ifis_training:  #Addnoiseforexplorationacrossexperts.  router_logits+=ntf.random_uniform(shape=router_logits.shape,minval=1-eps,maxval=1+eps)
+#Convertinputtosoftmaxoperationfrombfloati6float32forstability. router_logits=ntf.to_float32(router_logits)
+#Probabilitiesforeachtokenofwhatexpertistshouldbesentto. router_probs=ntf.softmax(router_logits,axis=-1)
+#Getthetop-1expertforeachtoken.expert_gateisthetop-1probability #fromtherouterforeachtoken.expert_indexiswhatexpertachtoken #isgoingtoberoutedto. #expert_gateshape:[num_cores,tokens_per_core] #expert_indexis:[num_cores,tokens_per_core] expert_gate,expert_index=ntf.top_1(router_probs,reduced_dim=num_experts)
+#expert_maskshape:[num_cores,tokens_per_core,num_experts] expert_mask=ntf.one_hot(expert_index,dimension=num_experts)
+#Computeloadbalancingloss. aux_loss=load_balance_loss(router_probs,expert_mask)
+#Expertshaveafixedcapacity,ensurewedoatececedit.Construct #thebatchindices,toeachexpert,withposition_in_expert #makewatthatmorethatexpert_capacitysamplescanberoutedto #eachexpert. position_in_expert=ntf.cumsum(expert_mask,dimension=tokens_per_core)*expert_mask
+#Keeponlytokensthatfitwithinexpert_capacity. expert_mask+=ntf.less(position_in_expert,expert_capacity) expert_mask_flat=ntf.reduce_sum(expert_mask,reduced_dim=experts_dim)
+#Maskouttheexpertshaveoverflow
+
+importmesh_tensorflowasntf defswitch_layer(inputs,n,capacity_factor,num_experts): """Distributedswitchtransformerfeed-forwardlayer.""" #num_cores(n)=totalcoresfortrainingthemodel(scalar). #d_model=modelhiddensize(scalar). #num_experts=totalnumberofexperts. #capacity_factor=extrabufferforeachexpert. #inputsshape:[batch,seq_len,d_model] batch,seq_len,d_model=inputs.get_shape() #Eachcorewillroutetokens_per_coretokenstothecorrectexperts. tokens_per_core=batch*seq_len/num_cores
+#Eachexpertwillhaveshape[num_cores,expert_capacity,d_model]. #Eachcoreisresponsibleforsendingexpert_capacitytokens
+#toeachexpert.expert_capacity=tokens_per_core*capacity_factor/num_experts
+#Reshapetosetuppercoreexpertdispatching. #shape:[batch,seq_len,d_model]->[num_cores,tokens_per_core,d_model] #Corelayout:[n,i,j]->[n,i,j]inputs=ntf.reshape(inputs,[num_cores,tokens_per_core,d_model])
+#Corelayout:[n,i,j]->[
+
+## References
+
+* Abadi et al. (2016) Martin Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: A system for large-scale machine learning. In _12th \(\{\)USENIX\(\}\) symposium on operating systems design and implementation (\(\{\)OSDI\(\}\) 16)_, pages 265-283, 2016.
+* Beltagy et al. (2020) Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. _arXiv preprint arXiv:2004.05150_, 2020.
+* Berant et al. (2013) Jonathan Berant, Andrew Chou, Roy Frostig, and Percy Liang. Semantic parsing on free-base from question-answer pairs. In _Proceedings of the 2013 conference on empirical methods in natural language processing_, pages 1533-1544, 2013.
+* Brown et al. (2020) Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. _arXiv preprint arXiv:2005.14165_, 2020.
+* Child et al. (2019) Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. _arXiv preprint arXiv:1904.10509_, 2019.
+* Cho and Bengio (2014) Kyunghyun Cho and Yoshua Bengio. Exponentially increasing the capacity-to-computation ratio for conditional computation in deep learning. _arXiv preprint arXiv:1406.7362_, 2014.
+* Clark et al. (2018) Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. _arXiv preprint arXiv:1803.05457_, 2018.
+* Correia et al. (2019) Goncalo M Correia, Vlad Niculae, and Andre FT Martins. Adaptively sparse transformers. _arXiv preprint arXiv:1909.00015_, 2019.
+* Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. _arXiv preprint arXiv:1810.04805_, 2018.
+* Eigen et al. (2013) David Eigen, Marc'Aurelio Ranzato, and Ilya Sutskever. Learning factored representations in a deep mixture of experts. _arXiv preprint arXiv:1312.4314_, 2013.
+* Fan et al. (2021) Angela Fan, Shruti Bhosale, Holger Schwenk, Zhiyi Ma, Ahmed El-Kishky, Siddharth Goyal, Mandeep Baines, Onur Celebi, Guillaume Wenzek, Vishrav Chaudhary, et al. Beyond english-centric multilingual machine translation. _Journal of Machine Learning Research_, 22(107):1-48, 2021.
+* Fedus et al. (2018) William Fedus, Ian Goodfellow, and Andrew M Dai. Maskgan: Better text generation via filling in the_. _arXiv preprint arXiv:1801.07736_, 2018.
+* Gale et al. (2020) Trevor Gale, Matei Zaharia, Cliff Young, and Erich Elsen. Sparse gpu kernels for deep learning. _arXiv preprint arXiv:2006.10901_, 2020.
+* Gray et al. (2017) Scott Gray, Alec Radford, and Diederik P Kingma. Gpu kernels for block-sparse weights. _[https://openai.com/blog/block-sparse-gpu-kernels/_](https://openai.com/blog/block-sparse-gpu-kernels/_), 2017.
+
+* Guu et al. (2020) Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-Wei Chang. Realm: Retrieval-augmented language model pre-training. _arXiv preprint arXiv:2002.08909_, 2020.
+* Harlap et al. (2018) Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Ganger, and Phil Gibbons. Pipedream: Fast and efficient pipeline parallel dnn training. _arXiv preprint arXiv:1806.03377_, 2018.
+* Hermann et al. (2015) Karl Moritz Hermann, Tomas Kocisky, Edward Grefenstette, Lasse Espeholt, Will Kay, Mustafa Suleyman, and Phil Blunsom. Teaching machines to read and comprehend. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett, editors, _Advances in Neural Information Processing Systems_, volume 28, pages 1693-1701. Curran Associates, Inc., 2015. URL [https://proceedings.neurips.cc/paper/2015/file/afdec7005cc9f14302cd0474fd0f3c96-Paper.pdf](https://proceedings.neurips.cc/paper/2015/file/afdec7005cc9f14302cd0474fd0f3c96-Paper.pdf).
+* Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. _arXiv preprint arXiv:1503.02531_, 2015.
+* Hochreiter and Schmidhuber (1997) Sepp Hochreiter and Jurgen Schmidhuber. Long short-term memory. _Neural computation_, 9(8):1735-1780, 1997.
+* Hooker (2020) Sara Hooker. The hardware lottery. _arXiv preprint arXiv:2009.06489_, 2020.
+* Huang et al. (2019) Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In _Advances in neural information processing systems_, pages 103-112, 2019.
+* Jacobs et al. (1991) Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. _Neural computation_, 3(1):79-87, 1991.
+* Jordan and Jacobs (1994) Michael I Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. _Neural computation_, 6(2):181-214, 1994.
+* Joshi et al. (2017) Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. _arXiv preprint arXiv:1705.03551_, 2017.
+* Kaplan et al. (2020) Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. _arXiv preprint arXiv:2001.08361_, 2020.
+* Kitaev et al. (2020) Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. _arXiv preprint arXiv:2001.04451_, 2020.
+* Kwiatkowski et al. (2019) Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: a benchmark for question answering research. _Transactions of the Association for Computational Linguistics_, 7:453-466, 2019.
+
+* Lample et al. (2019) Guillaume Lample, Alexandre Sablayrolles, Marc'Aurelio Ranzato, Ludovic Denoyer, and Herve Jegou. Large memory layers with product keys. In _Advances in Neural Information Processing Systems_, pages 8548-8559, 2019.
+* Lee et al. (2021) Katherine Lee, Daphne Ippolito, Andrew Nystrom, Chiyuan Zhang, Douglas Eck, Chris Callison-Burch, and Nicholas Carlini. Deduplicating training data makes language models better. _arXiv preprint arXiv:2107.06499_, 2021.
+* Lepikhin et al. (2020) Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. _arXiv preprint arXiv:2006.16668_, 2020.
+* Micikevicius et al. (2017) Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. _arXiv preprint arXiv:1710.03740_, 2017.
+* Narayan et al. (2018) Shashi Narayan, Shay B Cohen, and Mirella Lapata. Don't give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. _arXiv preprint arXiv:1808.08745_, 2018.
+* Nie et al. (2019) Yixin Nie, Adina Williams, Emily Dinan, Mohit Bansal, Jason Weston, and Douwe Kiela. Adversarial nli: A new benchmark for natural language understanding. _arXiv preprint arXiv:1910.14599_, 2019.
+* Puigcerver et al. (2020) Joan Puigcerver, Carlos Riquelme, Basil Mustafa, Cedric Renggli, Andre Susano Pinto, Sylvain Gelly, Daniel Keysers, and Neil Houlsby. Scalable transfer learning with expert models. _arXiv preprint arXiv:2009.13239_, 2020.
+* Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training, 2018.
+* Raffel et al. (2019) Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. _arXiv preprint arXiv:1910.10683_, 2019.
+* Rajbhandari et al. (2019) Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimization towards training a trillion parameter models. _arXiv preprint arXiv:1910.02054_, 2019.
+* Rajpurkar et al. (2016) Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. _arXiv preprint arXiv:1606.05250_, 2016.
+* Ramachandran and Le (2018) Prajit Ramachandran and Quoc V Le. Diversity and depth in per-example routing models. In _International Conference on Learning Representations_, 2018.
+* Robbins (1952) Herbert Robbins. Some aspects of the sequential design of experiments. _Bulletin of the American Mathematical Society_, 58(5):527-535, 1952.
+
+* Roberts et al. (2020) Adam Roberts, Colin Raffel, and Noam Shazeer. How much knowledge can you pack into the parameters of a language model? _arXiv preprint arXiv:2002.08910_, 2020.
+* Rosenbaum et al. (2017) Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. _arXiv preprint arXiv:1711.01239_, 2017.
+* Sakaguchi et al. (2020) Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In _Proceedings of the AAAI Conference on Artificial Intelligence_, volume 34, pages 8732-8740, 2020.
+* Sanh et al. (2019) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter, 2019.
+* Shazeer (2020) Noam Shazeer. Glu variants improve transformer, 2020.
+* Shazeer et al. (2017) Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. _arXiv preprint arXiv:1701.06538_, 2017.
+* Shazeer et al. (2018) Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, et al. Mesh-tensorflow: Deep learning for supercomputers. In _Advances in Neural Information Processing Systems_, pages 10414-10423, 2018.
+* Shoeybi et al. (2019) Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using gpu model parallelism. _arXiv preprint arXiv:1909.08053_, 2019.
+* Srivastava et al. (2014) Nitish Srivastava, Geoffrey E. Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. _Journal of Machine Learning Research_, 15(1):1929-1958, 2014. URL [http://www.cs.toronto.edu/~rsalakhu/papers/srivastava14a.pdf](http://www.cs.toronto.edu/~rsalakhu/papers/srivastava14a.pdf).
+* Strubell et al. (2019) Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in nlp. _arXiv preprint arXiv:1906.02243_, 2019.
+* Sukhbaatar et al. (2019) Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. _arXiv preprint arXiv:1905.07799_, 2019.
+* Sutton (2019) Rich Sutton. The Bitter Lesson. _[http://www.incompleteideas.net/IncIdeas/BitterLesson.html_](http://www.incompleteideas.net/IncIdeas/BitterLesson.html_), 2019.
+* Sutton and Barto (2018) Richard S Sutton and Andrew G Barto. _Reinforcement learning: An introduction_. Stanford University, 2018.
+* Taylor (1953) Wilson L Taylor. "cloze procedure": A new tool for measuring readability. _Journalism quarterly_, 30(4):415-433, 1953.
+
+* Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In _Advances in neural information processing systems_, pages 5998-6008, 2017.
+* Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. _arXiv preprint arXiv:1804.07461_, 2018.
+* Wang et al. (2019) Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. In _Advances in Neural Information Processing Systems_, pages 3266-3280, 2019.
+* Wang and Kanwar (2019) Shibo Wang and Pankaj Kanwar. Bfloat16: The secret to high performance on cloud tpus. _Google Cloud Blog_, 2019.
+* Xue et al. (2020) Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami Al-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mt5: A massively multilingual pre-trained text-to-text transformer. _arXiv preprint arXiv:2010.11934_, 2020.
+* Yang et al. (2020) Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. Xlnet: Generalized autoregressive pretraining for language understanding, 2020.
+* Zaheer et al. (2020) Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. _arXiv preprint arXiv:2007.14062_, 2020.

data/latex_to_md.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+
+# List all .tex files in the latex folder
+FILES=$(find latex -name "*.tex")
+
+for f in $FILES
+do
+  echo "Processing $f file..."
+  base_name=$(basename "$f" .tex)
+  out_file="references/${base_name}.md"
+
+ pandoc --wrap=none \
+         --no-highlight \
+         --strip-comments \
+         --from=latex \
+         --to=commonmark_x+pipe_tables \
+         "$f" \
+         -o "$out_file"
+  # Replace non-breaking spaces
+  sed -i .bak 's/ / /g' "$out_file"
+  sed -i .bak 's/ / /g' "$out_file"
+  sed -i .bak 's/ / /g' "$out_file"
+  sed -i .bak 's/ / /g' "$out_file"
+  sed -i.bak -E 's/`\\cite`//g; s/<[^>]*>//g; s/\{[^}]*\}//g; s/\\cite\{[^}]*\}//g' "$out_file"
+    sed -i.bak -E '
+    s/`\\cite`//g;   # Remove \cite commands inside backticks
+    s/::: //g;       # Remove the leading ::: for content markers
+    s/\[//g;         # Remove opening square bracket
+    s/\]//g;         # Remove closing square bracket
+  ' "$out_file"
+  # Remove .bak file
+  rm "$out_file.bak"
+done
+

docs/install_ocrmypdf.md

@@ -0,0 +1,29 @@
+## Linux
+
+- Run `apt-get install ocrmypdf`
+- Install ghostscript > 9.55 by following [these instructions](https://ghostscript.readthedocs.io/en/latest/Install.html) or running `scripts/install/ghostscript_install.sh`.
+- Run `pip install ocrmypdf`
+- Install any tesseract language packages that you want (example `apt-get install tesseract-ocr-eng`)
+- Set the tesseract data folder path
+  - Find the tesseract data folder `tessdata` with `find / -name tessdata`.  Make sure to use the one corresponding to the latest tesseract version if you have multiple.
+  - Create a `local.env` file in the root `marker` folder with `TESSDATA_PREFIX=/path/to/tessdata` inside it
+
+## Mac
+
+Only needed if using `ocrmypdf` as the ocr backend.
+
+- Run `brew install ocrmypdf`
+- Run `brew install tesseract-lang` to add language support
+- Run `pip install ocrmypdf`
+- Set the tesseract data folder path
+  - Find the tesseract data folder `tessdata` with `brew list tesseract`
+  - Create a `local.env` file in the root `marker` folder with `TESSDATA_PREFIX=/path/to/tessdata` inside it
+
+## Windows
+
+- Install `ocrmypdf` and ghostscript by following [these instructions](https://ocrmypdf.readthedocs.io/en/latest/installation.html#installing-on-windows)
+- Run `pip install ocrmypdf`
+- Install any tesseract language packages you want
+- Set the tesseract data folder path
+  - Find the tesseract data folder `tessdata` with `brew list tesseract`
+  - Create a `local.env` file in the root `marker` folder with `TESSDATA_PREFIX=/path/to/tessdata` inside it

marker/benchmark/scoring.py

@@ -0,0 +1,40 @@
+import math
+
+from rapidfuzz import fuzz
+import re
+import regex
+from statistics import mean
+
+CHUNK_MIN_CHARS = 25
+
+def chunk_text(text, chunk_len=500):
+    chunks = [text[i:i+chunk_len] for i in range(0, len(text), chunk_len)]
+    chunks = [c for c in chunks if c.strip() and len(c) > CHUNK_MIN_CHARS]
+    return chunks
+
+
+def overlap_score(hypothesis_chunks, reference_chunks):
+    length_modifier = len(hypothesis_chunks) / len(reference_chunks)
+    search_distance = max(len(reference_chunks) // 5, 10)
+    chunk_scores = []
+    for i, hyp_chunk in enumerate(hypothesis_chunks):
+        max_score = 0
+        total_len = 0
+        i_offset = int(i * length_modifier)
+        chunk_range = range(max(0, i_offset-search_distance), min(len(reference_chunks), i_offset+search_distance))
+        for j in chunk_range:
+            ref_chunk = reference_chunks[j]
+            score = fuzz.ratio(hyp_chunk, ref_chunk, score_cutoff=30) / 100
+            if score > max_score:
+                max_score = score
+                total_len = len(ref_chunk)
+        chunk_scores.append(max_score)
+    return chunk_scores
+
+
+def score_text(hypothesis, reference):
+    # Returns a 0-1 alignment score
+    hypothesis_chunks = chunk_text(hypothesis)
+    reference_chunks = chunk_text(reference)
+    chunk_scores = overlap_score(hypothesis_chunks, reference_chunks)
+    return mean(chunk_scores)

marker/cleaners/bullets.py

@@ -0,0 +1,8 @@
+import re
+
+
+def replace_bullets(text):
+    # Replace bullet characters with a -
+    bullet_pattern = r"(^|[\n ])[•●○■▪▫–—]( )"
+    replaced_string = re.sub(bullet_pattern, r"\1-\2", text)
+    return replaced_string

marker/cleaners/code.py

@@ -0,0 +1,131 @@
+from collections import Counter
+from statistics import mean, median
+
+from marker.schema.block import Span, Line
+from marker.schema.page import Page
+import re
+from typing import List
+
+
+def is_code_linelen(lines, thresh=80):
+    # Decide based on chars per newline threshold
+    total_alnum_chars = sum(len(re.findall(r'\w', line.prelim_text)) for line in lines)
+    total_newlines = max(len(lines) - 1, 1)
+
+    if total_alnum_chars == 0:
+        return False
+
+    ratio = total_alnum_chars / total_newlines
+    return ratio < thresh
+
+
+def comment_count(lines):
+    pattern = re.compile(r"^(//|#|'|--|/\*|'''|\"\"\"|--\[\[|<!--|%|%{|\(\*)")
+    return sum([1 for line in lines if pattern.match(line)])
+
+
+def identify_code_blocks(pages: List[Page]):
+    code_block_count = 0
+    font_sizes = []
+    line_heights = []
+    for page in pages:
+        font_sizes += page.get_font_sizes()
+        line_heights += page.get_line_heights()
+
+    avg_font_size = None
+    avg_line_height = None
+    if len(font_sizes) > 0:
+        avg_line_height = median(line_heights)
+        avg_font_size = mean(font_sizes)
+
+    for page in pages:
+        for block in page.blocks:
+            if block.block_type != "Text":
+                last_block = block
+                continue
+
+            # Ensure we have lines and spans
+            if len(block.lines) == 0:
+                continue
+            if sum([len(line.spans) for line in block.lines]) == 0:
+                continue
+
+            min_start = block.get_min_line_start()
+
+            is_indent = []
+            line_fonts = []
+            line_font_sizes = []
+            block_line_heights = []
+            for line in block.lines:
+                line_fonts += [span.font for span in line.spans]
+                line_font_sizes += [span.font_size for span in line.spans]
+                block_line_heights.append(line.bbox[3] - line.bbox[1])
+
+                is_indent.append(line.bbox[0] > min_start)
+
+            comment_lines = comment_count([line.prelim_text for line in block.lines])
+            is_code = [
+                len(block.lines) > 3,
+                is_code_linelen(block.lines),
+                sum(is_indent) + comment_lines > len(block.lines) * .7, # Indentation and comments are a majority
+            ]
+
+            if avg_font_size is not None:
+                font_checks = [
+                    mean(line_font_sizes) <= avg_font_size * .8, # Lower than average font size and line height
+                    mean(block_line_heights) < avg_line_height * .8
+                ]
+                is_code += font_checks
+
+            if all(is_code):
+                code_block_count += 1
+                block.block_type = "Code"
+
+    return code_block_count
+
+
+def indent_blocks(pages: List[Page]):
+    span_counter = 0
+    for page in pages:
+        for block in page.blocks:
+            if block.block_type != "Code":
+                continue
+
+            lines = []
+            min_left = 1000  # will contain x- coord of column 0
+            col_width = 0  # width of 1 char
+            for line in block.lines:
+                text = ""
+                min_left = min(line.bbox[0], min_left)
+                for span in line.spans:
+                    if col_width == 0 and len(span.text) > 0:
+                        col_width = (span.bbox[2] - span.bbox[0]) / len(span.text)
+                    text += span.text
+                lines.append((line.bbox, text))
+
+            block_text = ""
+            blank_line = False
+            for line in lines:
+                text = line[1]
+                if col_width == 0:
+                    prefix = ""
+                else:
+                    prefix = " " * int((line[0][0] - min_left) / col_width)
+                current_line_blank = len(text.strip()) == 0
+                if blank_line and current_line_blank:
+                    # Don't put multiple blank lines in a row
+                    continue
+
+                block_text += prefix + text + "\n"
+                blank_line = current_line_blank
+
+            new_span = Span(
+                text=block_text,
+                bbox=block.bbox,
+                span_id=f"{span_counter}_fix_code",
+                font=block.lines[0].spans[0].font,
+                font_weight=block.lines[0].spans[0].font_weight,
+                font_size=block.lines[0].spans[0].font_size,
+            )
+            span_counter += 1
+            block.lines = [Line(spans=[new_span], bbox=block.bbox)]

marker/cleaners/fontstyle.py

@@ -0,0 +1,30 @@
+from typing import List
+
+from marker.schema.page import Page
+
+
+def find_bold_italic(pages: List[Page], bold_min_weight=600):
+    font_weights = []
+    for page in pages:
+        for block in page.blocks:
+            # We don't want to bias our font stats
+            if block.block_type in ["Title", "Section-header"]:
+                continue
+            for line in block.lines:
+                for span in line.spans:
+                    if "bold" in span.font.lower():
+                        span.bold = True
+                    if "ital" in span.font.lower():
+                        span.italic = True
+
+                    font_weights.append(span.font_weight)
+
+    if len(font_weights) == 0:
+        return
+
+    for page in pages:
+        for block in page.blocks:
+            for line in block.lines:
+                for span in line.spans:
+                    if span.font_weight >= bold_min_weight:
+                        span.bold = True

marker/cleaners/headers.py

@@ -0,0 +1,82 @@
+import re
+from collections import Counter
+from rapidfuzz import fuzz
+
+from marker.schema.merged import FullyMergedBlock
+from typing import List, Tuple
+
+
+def filter_common_elements(lines, page_count, threshold=.6):
+    # We can't filter if we don't have enough pages to find common elements
+    if page_count < 3:
+        return []
+    text = [s.text for line in lines for s in line.spans if len(s.text) > 4]
+    counter = Counter(text)
+    common = [k for k, v in counter.items() if v > page_count * threshold]
+    bad_span_ids = [s.span_id for line in lines for s in line.spans if s.text in common]
+    return bad_span_ids
+
+
+def filter_header_footer(all_page_blocks, max_selected_lines=2):
+    first_lines = []
+    last_lines = []
+    for page in all_page_blocks:
+        nonblank_lines = page.get_nonblank_lines()
+        first_lines.extend(nonblank_lines[:max_selected_lines])
+        last_lines.extend(nonblank_lines[-max_selected_lines:])
+
+    bad_span_ids = filter_common_elements(first_lines, len(all_page_blocks))
+    bad_span_ids += filter_common_elements(last_lines, len(all_page_blocks))
+    return bad_span_ids
+
+
+def replace_leading_trailing_digits(string, replacement):
+    string = re.sub(r'^\d+', replacement, string)
+    string = re.sub(r'\d+$', replacement, string)
+    return string
+
+
+def find_overlap_elements(lst: List[Tuple[str, int]], string_match_thresh=.9, min_overlap=.05) -> List[int]:
+    # Initialize a list to store the elements that meet the criteria
+    result = []
+    titles = [l[0] for l in lst]
+
+    for i, (str1, id_num) in enumerate(lst):
+        overlap_count = 0  # Count the number of elements that overlap by at least 80%
+
+        for j, str2 in enumerate(titles):
+            if i != j and fuzz.ratio(str1, str2) >= string_match_thresh * 100:
+                overlap_count += 1
+
+        # Check if the element overlaps with at least 50% of other elements
+        if overlap_count >= max(3.0, len(lst) * min_overlap):
+            result.append(id_num)
+
+    return result
+
+
+def filter_common_titles(merged_blocks: List[FullyMergedBlock]) -> List[FullyMergedBlock]:
+    titles = []
+    for i, block in enumerate(merged_blocks):
+        if block.block_type in ["Title", "Section-header"]:
+            text = block.text
+            if text.strip().startswith("#"):
+                text = re.sub(r'#+', '', text)
+            text = text.strip()
+            # Remove page numbers from start/end
+            text = replace_leading_trailing_digits(text, "").strip()
+            titles.append((text, i))
+
+    bad_block_ids = find_overlap_elements(titles)
+
+    new_blocks = []
+    for i, block in enumerate(merged_blocks):
+        if i in bad_block_ids:
+            continue
+        new_blocks.append(block)
+
+    return new_blocks
+
+
+
+

marker/cleaners/headings.py

@@ -0,0 +1,59 @@
+from typing import List
+
+from marker.settings import settings
+from marker.schema.bbox import rescale_bbox
+from marker.schema.block import bbox_from_lines
+from marker.schema.page import Page
+
+
+def split_heading_blocks(pages: List[Page]):
+    # Heading lines can be combined into regular text blocks sometimes by pdftext
+    # Split up heading lines into separate blocks properly
+    for page in pages:
+        page_heading_boxes = [b for b in page.layout.bboxes if b.label in ["Title", "Section-header"]]
+        page_heading_boxes = [(rescale_bbox(page.layout.image_bbox, page.bbox, b.bbox), b.label) for b in page_heading_boxes]
+
+        new_blocks = []
+        for block_idx, block in enumerate(page.blocks):
+            if block.block_type not in ["Text"]:
+                new_blocks.append(block)
+                continue
+
+            heading_lines = []
+            for line_idx, line in enumerate(block.lines):
+                for (heading_box, label) in page_heading_boxes:
+                    if line.intersection_pct(heading_box) > settings.BBOX_INTERSECTION_THRESH:
+                        heading_lines.append((line_idx, label))
+                        break
+
+            if len(heading_lines) == 0:
+                new_blocks.append(block)
+                continue
+
+            # Split up the block into separate blocks around headers
+            start = 0
+            for (heading_line, label) in heading_lines:
+                if start < heading_line:
+                    copied_block = block.copy()
+                    copied_block.lines = block.lines[start:heading_line]
+                    copied_block.bbox = bbox_from_lines(copied_block.lines)
+                    new_blocks.append(copied_block)
+
+                copied_block = block.copy()
+                copied_block.lines = block.lines[heading_line:heading_line + 1]
+                copied_block.block_type = label
+                copied_block.bbox = bbox_from_lines(copied_block.lines)
+                new_blocks.append(copied_block)
+
+                start = heading_line + 1
+                if start >= len(block.lines):
+                    break
+
+            # Add any remaining lines
+            if start < len(block.lines):
+                copied_block = block.copy()
+                copied_block.lines = block.lines[start:]
+                copied_block.bbox = bbox_from_lines(copied_block.lines)
+                new_blocks.append(copied_block)
+
+        page.blocks = new_blocks

marker/cleaners/text.py

@@ -0,0 +1,8 @@
+import re
+
+
+def cleanup_text(full_text):
+    full_text = re.sub(r'\n{3,}', '\n\n', full_text)
+    full_text = re.sub(r'(\n\s){3,}', '\n\n', full_text)
+    full_text = full_text.replace('\xa0', ' ') # Replace non-breaking spaces
+    return full_text

marker/convert.py

@@ -0,0 +1,162 @@
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning) # Filter torch pytree user warnings
+
+import pypdfium2 as pdfium
+from PIL import Image
+
+from marker.utils import flush_cuda_memory
+from marker.tables.table import format_tables
+from marker.debug.data import dump_bbox_debug_data
+from marker.layout.layout import surya_layout, annotate_block_types
+from marker.layout.order import surya_order, sort_blocks_in_reading_order
+from marker.ocr.lang import replace_langs_with_codes, validate_langs
+from marker.ocr.detection import surya_detection
+from marker.ocr.recognition import run_ocr
+from marker.pdf.extract_text import get_text_blocks
+from marker.cleaners.headers import filter_header_footer, filter_common_titles
+from marker.equations.equations import replace_equations
+from marker.pdf.utils import find_filetype
+from marker.postprocessors.editor import edit_full_text
+from marker.cleaners.code import identify_code_blocks, indent_blocks
+from marker.cleaners.bullets import replace_bullets
+from marker.cleaners.headings import split_heading_blocks
+from marker.cleaners.fontstyle import find_bold_italic
+from marker.postprocessors.markdown import merge_spans, merge_lines, get_full_text
+from marker.cleaners.text import cleanup_text
+from marker.images.extract import extract_images
+from marker.images.save import images_to_dict
+
+from typing import List, Dict, Tuple, Optional
+from marker.settings import settings
+
+
+def convert_single_pdf(
+        fname: str,
+        model_lst: List,
+        max_pages=None,
+        metadata: Optional[Dict]=None,
+        langs: Optional[List[str]] = None,
+        batch_multiplier: int = 1
+) -> Tuple[str, Dict[str, Image.Image], Dict]:
+    # Set language needed for OCR
+    if langs is None:
+        langs = [settings.DEFAULT_LANG]
+
+    if metadata:
+        langs = metadata.get("languages", langs)
+
+    langs = replace_langs_with_codes(langs)
+    validate_langs(langs)
+
+    # Find the filetype
+    filetype = find_filetype(fname)
+
+    # Setup output metadata
+    out_meta = {
+        "languages": langs,
+        "filetype": filetype,
+    }
+
+    if filetype == "other": # We can't process this file
+        return "", out_meta
+
+    # Get initial text blocks from the pdf
+    doc = pdfium.PdfDocument(fname)
+    pages, toc = get_text_blocks(
+        doc,
+        max_pages=max_pages,
+    )
+    out_meta.update({
+        "toc": toc,
+        "pages": len(pages),
+    })
+
+    # Unpack models from list
+    texify_model, layout_model, order_model, edit_model, detection_model, ocr_model = model_lst
+
+    # Identify text lines on pages
+    surya_detection(doc, pages, detection_model, batch_multiplier=batch_multiplier)
+    flush_cuda_memory()
+
+    # OCR pages as needed
+    pages, ocr_stats = run_ocr(doc, pages, langs, ocr_model, batch_multiplier=batch_multiplier)
+    flush_cuda_memory()
+
+    out_meta["ocr_stats"] = ocr_stats
+    if len([b for p in pages for b in p.blocks]) == 0:
+        print(f"Could not extract any text blocks for {fname}")
+        return "", out_meta
+
+    surya_layout(doc, pages, layout_model, batch_multiplier=batch_multiplier)
+    flush_cuda_memory()
+
+    # Find headers and footers
+    bad_span_ids = filter_header_footer(pages)
+    out_meta["block_stats"] = {"header_footer": len(bad_span_ids)}
+
+    # Add block types in
+    annotate_block_types(pages)
+
+    # Dump debug data if flags are set
+    dump_bbox_debug_data(doc, pages)
+
+    # Find reading order for blocks
+    # Sort blocks by reading order
+    surya_order(doc, pages, order_model, batch_multiplier=batch_multiplier)
+    sort_blocks_in_reading_order(pages)
+    flush_cuda_memory()
+
+    # Fix code blocks
+    code_block_count = identify_code_blocks(pages)
+    out_meta["block_stats"]["code"] = code_block_count
+    indent_blocks(pages)
+
+    # Fix table blocks
+    table_count = format_tables(pages)
+    out_meta["block_stats"]["table"] = table_count
+
+    for page in pages:
+        for block in page.blocks:
+            block.filter_spans(bad_span_ids)
+            block.filter_bad_span_types()
+
+    filtered, eq_stats = replace_equations(
+        doc,
+        pages,
+        texify_model,
+        batch_multiplier=batch_multiplier
+    )
+    flush_cuda_memory()
+    out_meta["block_stats"]["equations"] = eq_stats
+
+    # Extract images and figures
+    if settings.EXTRACT_IMAGES:
+        extract_images(doc, pages)
+
+    # Split out headers
+    split_heading_blocks(pages)
+    find_bold_italic(pages)
+
+    # Copy to avoid changing original data
+    merged_lines = merge_spans(filtered)
+    text_blocks = merge_lines(merged_lines)
+    text_blocks = filter_common_titles(text_blocks)
+    full_text = get_full_text(text_blocks)
+
+    # Handle empty blocks being joined
+    full_text = cleanup_text(full_text)
+
+    # Replace bullet characters with a -
+    full_text = replace_bullets(full_text)
+
+    # Postprocess text with editor model
+    full_text, edit_stats = edit_full_text(
+        full_text,
+        edit_model,
+        batch_multiplier=batch_multiplier
+    )
+    flush_cuda_memory()
+    out_meta["postprocess_stats"] = {"edit": edit_stats}
+    doc_images = images_to_dict(pages)
+
+    return full_text, doc_images, out_meta

marker/debug/data.py

@@ -0,0 +1,76 @@
+import base64
+import json
+import os
+from typing import List
+
+from marker.pdf.images import render_image
+from marker.schema.page import Page
+from marker.settings import settings
+from PIL import Image
+import io
+
+
+def dump_equation_debug_data(doc, images, converted_spans):
+    if not settings.DEBUG_DATA_FOLDER or settings.DEBUG_LEVEL == 0:
+        return
+
+    if len(images) == 0:
+        return
+
+    # We attempted one conversion per image
+    assert len(converted_spans) == len(images)
+
+    data_lines = []
+    for idx, (pil_image, converted_span) in enumerate(zip(images, converted_spans)):
+        if converted_span is None:
+            continue
+        # Image is a BytesIO object
+        img_bytes = io.BytesIO()
+        pil_image.save(img_bytes, format="WEBP", lossless=True)
+        b64_image = base64.b64encode(img_bytes.getvalue()).decode("utf-8")
+        data_lines.append({
+            "image": b64_image,
+            "text": converted_span.text,
+            "bbox": converted_span.bbox
+        })
+
+    # Remove extension from doc name
+    doc_base = os.path.basename(doc.name).rsplit(".", 1)[0]
+
+    debug_file = os.path.join(settings.DEBUG_DATA_FOLDER, f"{doc_base}_equations.json")
+    with open(debug_file, "w+") as f:
+        json.dump(data_lines, f)
+
+
+def dump_bbox_debug_data(doc, blocks: List[Page]):
+    if not settings.DEBUG_DATA_FOLDER or settings.DEBUG_LEVEL < 2:
+        return
+
+    # Remove extension from doc name
+    doc_base = os.path.basename(doc.name).rsplit(".", 1)[0]
+
+    debug_file = os.path.join(settings.DEBUG_DATA_FOLDER, f"{doc_base}_bbox.json")
+    debug_data = []
+    for idx, page_blocks in enumerate(blocks):
+        page = doc[idx]
+
+        png_image = render_image(page, dpi=settings.TEXIFY_DPI)
+        width, height = png_image.size
+        max_dimension = 6000
+        if width > max_dimension or height > max_dimension:
+            scaling_factor = min(max_dimension / width, max_dimension / height)
+            png_image = png_image.resize((int(width * scaling_factor), int(height * scaling_factor)), Image.ANTIALIAS)
+
+        img_bytes = io.BytesIO()
+        png_image.save(img_bytes, format="WEBP", lossless=True, quality=100)
+        b64_image = base64.b64encode(img_bytes.getvalue()).decode("utf-8")
+
+        page_data = page_blocks.model_dump()
+        page_data["image"] = b64_image
+        debug_data.append(page_data)
+
+    with open(debug_file, "w+") as f:
+        json.dump(debug_data, f)
+
+
+

marker/equations/equations.py

@@ -0,0 +1,183 @@
+from collections import defaultdict
+from copy import deepcopy
+from typing import List
+
+from marker.debug.data import dump_equation_debug_data
+from marker.equations.inference import get_total_texify_tokens, get_latex_batched
+from marker.pdf.images import render_bbox_image
+from marker.schema.bbox import rescale_bbox
+from marker.schema.page import Page
+from marker.schema.block import Line, Span, Block, bbox_from_lines, split_block_lines, find_insert_block
+from marker.settings import settings
+
+
+def find_equation_blocks(page, processor):
+    equation_blocks = []
+    equation_regions = [l.bbox for l in page.layout.bboxes if l.label in ["Formula"]]
+    equation_regions = [rescale_bbox(page.layout.image_bbox, page.bbox, b) for b in equation_regions]
+
+    lines_to_remove = defaultdict(list)
+    insert_points = {}
+    equation_lines = defaultdict(list)
+    for region_idx, region in enumerate(equation_regions):
+        for block_idx, block in enumerate(page.blocks):
+            for line_idx, line in enumerate(block.lines):
+                if line.intersection_pct(region) > settings.BBOX_INTERSECTION_THRESH:
+                    # We will remove this line from the block
+                    lines_to_remove[region_idx].append((block_idx, line_idx))
+                    equation_lines[region_idx].append(line)
+
+                    if region_idx not in insert_points:
+                        insert_points[region_idx] = (block_idx, line_idx)
+
+    # Account for regions where the lines were not detected
+    for region_idx, region in enumerate(equation_regions):
+        if region_idx in insert_points:
+            continue
+
+        insert_points[region_idx] = (find_insert_block(page.blocks, region), 0)
+
+    block_lines_to_remove = defaultdict(set)
+    for region_idx, equation_region in enumerate(equation_regions):
+        if region_idx not in equation_lines or len(equation_lines[region_idx]) == 0:
+            block_text = ""
+            total_tokens = 0
+        else:
+            equation_block = equation_lines[region_idx]
+            block_text = " ".join([line.prelim_text for line in equation_block])
+            total_tokens = get_total_texify_tokens(block_text, processor)
+
+        equation_insert = insert_points[region_idx]
+        equation_insert_line_idx = equation_insert[1]
+        equation_insert_line_idx -= len(
+            [x for x in lines_to_remove[region_idx] if x[0] == equation_insert[0] and x[1] < equation_insert[1]])
+
+        selected_blocks = [equation_insert[0], equation_insert_line_idx, total_tokens, block_text, equation_region]
+        if total_tokens < settings.TEXIFY_MODEL_MAX:
+            # Account for the lines we're about to remove
+            for item in lines_to_remove[region_idx]:
+                block_lines_to_remove[item[0]].add(item[1])
+            equation_blocks.append(selected_blocks)
+
+    # Remove the lines from the blocks
+    for block_idx, bad_lines in block_lines_to_remove.items():
+        block = page.blocks[block_idx]
+        block.lines = [line for idx, line in enumerate(block.lines) if idx not in bad_lines]
+
+    return equation_blocks
+
+
+def increment_insert_points(page_equation_blocks, insert_block_idx, insert_count):
+    for idx, (block_idx, line_idx, token_count, block_text, equation_bbox) in enumerate(page_equation_blocks):
+        if block_idx >= insert_block_idx:
+            page_equation_blocks[idx][0] += insert_count
+
+
+def insert_latex_block(page_blocks: Page, page_equation_blocks, predictions, pnum, processor):
+    converted_spans = []
+    idx = 0
+    success_count = 0
+    fail_count = 0
+    for block_number, (insert_block_idx, insert_line_idx, token_count, block_text, equation_bbox) in enumerate(page_equation_blocks):
+        latex_text = predictions[block_number]
+        conditions = [
+            get_total_texify_tokens(latex_text, processor) < settings.TEXIFY_MODEL_MAX,  # Make sure we didn't get to the overall token max, indicates run-on
+            len(latex_text) > len(block_text) * .7,
+            len(latex_text.strip()) > 0
+        ]
+
+        new_block = Block(
+            lines=[Line(
+                spans=[
+                    Span(
+                        text=block_text.replace("\n", " "),
+                        bbox=equation_bbox,
+                        span_id=f"{pnum}_{idx}_fixeq",
+                        font="Latex",
+                        font_weight=0,
+                        font_size=0
+                    )
+                ],
+                bbox=equation_bbox
+            )],
+            bbox=equation_bbox,
+            block_type="Formula",
+            pnum=pnum
+        )
+
+        if not all(conditions):
+            fail_count += 1
+        else:
+            success_count += 1
+            new_block.lines[0].spans[0].text = latex_text.replace("\n", " ")
+            converted_spans.append(deepcopy(new_block.lines[0].spans[0]))
+
+        # Add in the new LaTeX block
+        if insert_line_idx == 0:
+            page_blocks.blocks.insert(insert_block_idx, new_block)
+            increment_insert_points(page_equation_blocks, insert_block_idx, 1)
+        elif insert_line_idx >= len(page_blocks.blocks[insert_block_idx].lines):
+            page_blocks.blocks.insert(insert_block_idx + 1, new_block)
+            increment_insert_points(page_equation_blocks, insert_block_idx + 1, 1)
+        else:
+            new_blocks = []
+            for block_idx, block in enumerate(page_blocks.blocks):
+                if block_idx == insert_block_idx:
+                    split_block = split_block_lines(block, insert_line_idx)
+                    new_blocks.append(split_block[0])
+                    new_blocks.append(new_block)
+                    new_blocks.append(split_block[1])
+                    increment_insert_points(page_equation_blocks, insert_block_idx, 2)
+                else:
+                    new_blocks.append(block)
+            page_blocks.blocks = new_blocks
+
+    return success_count, fail_count, converted_spans
+
+
+def replace_equations(doc, pages: List[Page], texify_model, batch_multiplier=1):
+    unsuccessful_ocr = 0
+    successful_ocr = 0
+
+    # Find potential equation regions, and length of text in each region
+    equation_blocks = []
+    for pnum, page in enumerate(pages):
+        equation_blocks.append(find_equation_blocks(page, texify_model.processor))
+
+    eq_count = sum([len(x) for x in equation_blocks])
+
+    images = []
+    token_counts = []
+    for page_idx, page_equation_blocks in enumerate(equation_blocks):
+        page_obj = doc[page_idx]
+        for equation_idx, (insert_block_idx, insert_line_idx, token_count, block_text, equation_bbox) in enumerate(page_equation_blocks):
+            png_image = render_bbox_image(page_obj, pages[page_idx], equation_bbox)
+
+            images.append(png_image)
+            token_counts.append(token_count)
+
+    # Make batched predictions
+    predictions = get_latex_batched(images, token_counts, texify_model, batch_multiplier=batch_multiplier)
+
+    # Replace blocks with predictions
+    page_start = 0
+    converted_spans = []
+    for page_idx, page_equation_blocks in enumerate(equation_blocks):
+        page_equation_count = len(page_equation_blocks)
+        page_predictions = predictions[page_start:page_start + page_equation_count]
+        success_count, fail_count, converted_span = insert_latex_block(
+            pages[page_idx],
+            page_equation_blocks,
+            page_predictions,
+            page_idx,
+            texify_model.processor
+        )
+        converted_spans.extend(converted_span)
+        page_start += page_equation_count
+        successful_ocr += success_count
+        unsuccessful_ocr += fail_count
+
+    # If debug mode is on, dump out conversions for comparison
+    dump_equation_debug_data(doc, images, converted_spans)
+
+    return pages, {"successful_ocr": successful_ocr, "unsuccessful_ocr": unsuccessful_ocr, "equations": eq_count}

marker/equations/inference.py

@@ -0,0 +1,50 @@
+from texify.inference import batch_inference
+
+from marker.settings import settings
+import os
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+def get_batch_size():
+    if settings.TEXIFY_BATCH_SIZE is not None:
+        return settings.TEXIFY_BATCH_SIZE
+    elif settings.TORCH_DEVICE_MODEL == "cuda":
+        return 6
+    elif settings.TORCH_DEVICE_MODEL == "mps":
+        return 6
+    return 2
+
+def get_latex_batched(images, token_counts, texify_model, batch_multiplier=1):
+    if len(images) == 0:
+        return []
+
+    predictions = [""] * len(images)
+    batch_size = get_batch_size() * batch_multiplier
+
+    for i in range(0, len(images), batch_size):
+        # Dynamically set max length to save inference time
+        min_idx = i
+        max_idx = min(min_idx + batch_size, len(images))
+        max_length = max(token_counts[min_idx:max_idx])
+        max_length = min(max_length, settings.TEXIFY_MODEL_MAX)
+        max_length += settings.TEXIFY_TOKEN_BUFFER
+
+        model_output = batch_inference(images[min_idx:max_idx], texify_model, texify_model.processor, max_tokens=max_length)
+
+        for j, output in enumerate(model_output):
+            token_count = get_total_texify_tokens(output, texify_model.processor)
+            if token_count >= max_length - 1:
+                output = ""
+
+            image_idx = i + j
+            predictions[image_idx] = output
+    return predictions
+
+
+def get_total_texify_tokens(text, processor):
+    tokenizer = processor.tokenizer
+    tokens = tokenizer(text)
+    return len(tokens["input_ids"])
+
+

marker/images/extract.py

@@ -0,0 +1,72 @@
+from marker.images.save import get_image_filename
+from marker.pdf.images import render_bbox_image
+from marker.schema.bbox import rescale_bbox
+from marker.schema.block import find_insert_block, Span, Line
+from marker.settings import settings
+
+
+def find_image_blocks(page):
+    image_blocks = []
+    image_regions = [l.bbox for l in page.layout.bboxes if l.label in ["Figure", "Picture"]]
+    image_regions = [rescale_bbox(page.layout.image_bbox, page.bbox, b) for b in image_regions]
+
+    insert_points = {}
+    for region_idx, region in enumerate(image_regions):
+        for block_idx, block in enumerate(page.blocks):
+            for line_idx, line in enumerate(block.lines):
+                if line.intersection_pct(region) > settings.BBOX_INTERSECTION_THRESH:
+                    line.spans = [] # We will remove this line from the block
+
+                    if region_idx not in insert_points:
+                        insert_points[region_idx] = (block_idx, line_idx)
+
+    # Account for images with no detected lines
+    for region_idx, region in enumerate(image_regions):
+        if region_idx in insert_points:
+            continue
+
+        insert_points[region_idx] = (find_insert_block(page.blocks, region), 0)
+
+    for region_idx, image_region in enumerate(image_regions):
+        image_insert = insert_points[region_idx]
+        image_blocks.append([image_insert[0], image_insert[1], image_region])
+
+    return image_blocks
+
+
+def extract_page_images(page_obj, page):
+    page.images = []
+    image_blocks = find_image_blocks(page)
+
+    for image_idx, (block_idx, line_idx, bbox) in enumerate(image_blocks):
+        block = page.blocks[block_idx]
+        image = render_bbox_image(page_obj, page, bbox)
+        image_filename = get_image_filename(page, image_idx)
+        image_markdown = f"\n\n![{image_filename}]({image_filename})\n\n"
+        image_span = Span(
+            bbox=bbox,
+            text=image_markdown,
+            font="Image",
+            rotation=0,
+            font_weight=0,
+            font_size=0,
+            image=True,
+            span_id=f"image_{image_idx}"
+        )
+
+        # Sometimes, the block has zero lines
+        if len(block.lines) > line_idx:
+            block.lines[line_idx].spans.append(image_span)
+        else:
+            line = Line(
+                bbox=bbox,
+                spans=[image_span]
+            )
+            block.lines.append(line)
+        page.images.append(image)
+
+
+def extract_images(doc, pages):
+    for page_idx, page in enumerate(pages):
+        page_obj = doc[page_idx]
+        extract_page_images(page_obj, page)

marker/images/save.py

@@ -0,0 +1,18 @@
+from typing import List
+
+from marker.schema.page import Page
+
+
+def get_image_filename(page: Page, image_idx):
+    return f"{page.pnum}_image_{image_idx}.png"
+
+
+def images_to_dict(pages: List[Page]):
+    images = {}
+    for page in pages:
+        if page.images is None:
+            continue
+        for image_idx, image in enumerate(page.images):
+            image_filename = get_image_filename(page, image_idx)
+            images[image_filename] = image
+    return images

marker/layout/layout.py

@@ -0,0 +1,48 @@
+from typing import List
+
+from surya.layout import batch_layout_detection
+
+from marker.pdf.images import render_image
+from marker.schema.bbox import rescale_bbox
+from marker.schema.page import Page
+from marker.settings import settings
+
+
+def get_batch_size():
+    if settings.LAYOUT_BATCH_SIZE is not None:
+        return settings.LAYOUT_BATCH_SIZE
+    elif settings.TORCH_DEVICE_MODEL == "cuda":
+        return 6
+    return 6
+
+
+def surya_layout(doc, pages: List[Page], layout_model, batch_multiplier=1):
+    images = [render_image(doc[pnum], dpi=settings.SURYA_LAYOUT_DPI) for pnum in range(len(pages))]
+    text_detection_results = [p.text_lines for p in pages]
+
+    processor = layout_model.processor
+    layout_results = batch_layout_detection(images, layout_model, processor, detection_results=text_detection_results, batch_size=get_batch_size() * batch_multiplier)
+    for page, layout_result in zip(pages, layout_results):
+        page.layout = layout_result
+
+
+def annotate_block_types(pages: List[Page]):
+    for page in pages:
+        max_intersections = {}
+        for i, block in enumerate(page.blocks):
+            for j, layout_block in enumerate(page.layout.bboxes):
+                layout_bbox = layout_block.bbox
+                layout_bbox = rescale_bbox(page.layout.image_bbox, page.bbox, layout_bbox)
+                intersection_pct = block.intersection_pct(layout_bbox)
+                if i not in max_intersections:
+                    max_intersections[i] = (intersection_pct, j)
+                elif intersection_pct > max_intersections[i][0]:
+                    max_intersections[i] = (intersection_pct, j)
+
+        for i, block in enumerate(page.blocks):
+            block = page.blocks[i]
+            block_type = "Text"
+            if i in max_intersections:
+                j = max_intersections[i][1]
+                block_type = page.layout.bboxes[j].label
+            block.block_type = block_type

marker/layout/order.py

@@ -0,0 +1,69 @@
+from collections import defaultdict
+from typing import List
+
+from surya.ordering import batch_ordering
+
+from marker.pdf.images import render_image
+from marker.pdf.utils import sort_block_group
+from marker.schema.bbox import rescale_bbox
+from marker.schema.page import Page
+from marker.settings import settings
+
+
+def get_batch_size():
+    if settings.ORDER_BATCH_SIZE is not None:
+        return settings.ORDER_BATCH_SIZE
+    elif settings.TORCH_DEVICE_MODEL == "cuda":
+        return 6
+    elif settings.TORCH_DEVICE_MODEL == "mps":
+        return 6
+    return 6
+
+
+def surya_order(doc, pages: List[Page], order_model, batch_multiplier=1):
+    images = [render_image(doc[pnum], dpi=settings.SURYA_ORDER_DPI) for pnum in range(len(pages))]
+
+    # Get bboxes for all pages
+    bboxes = []
+    for page in pages:
+        bbox = [b.bbox for b in page.layout.bboxes][:settings.ORDER_MAX_BBOXES]
+        bboxes.append(bbox)
+
+    processor = order_model.processor
+    order_results = batch_ordering(images, bboxes, order_model, processor, batch_size=get_batch_size() * batch_multiplier)
+    for page, order_result in zip(pages, order_results):
+        page.order = order_result
+
+
+def sort_blocks_in_reading_order(pages: List[Page]):
+    for page in pages:
+        order = page.order
+        block_positions = {}
+        max_position = 0
+        for i, block in enumerate(page.blocks):
+            for order_box in order.bboxes:
+                order_bbox = order_box.bbox
+                position = order_box.position
+                order_bbox = rescale_bbox(order.image_bbox, page.bbox, order_bbox)
+                block_intersection = block.intersection_pct(order_bbox)
+                if i not in block_positions:
+                    block_positions[i] = (block_intersection, position)
+                elif block_intersection > block_positions[i][0]:
+                    block_positions[i] = (block_intersection, position)
+                max_position = max(max_position, position)
+        block_groups = defaultdict(list)
+        for i, block in enumerate(page.blocks):
+            if i in block_positions:
+                position = block_positions[i][1]
+            else:
+                max_position += 1
+                position = max_position
+
+            block_groups[position].append(block)
+
+        new_blocks = []
+        for position in sorted(block_groups.keys()):
+            block_group = sort_block_group(block_groups[position])
+            new_blocks.extend(block_group)
+
+        page.blocks = new_blocks

marker/logger.py

@@ -0,0 +1,12 @@
+import logging
+import warnings
+
+
+def configure_logging():
+    logging.basicConfig(level=logging.WARNING)
+
+    logging.getLogger('pdfminer').setLevel(logging.ERROR)
+    logging.getLogger('PIL').setLevel(logging.ERROR)
+    logging.getLogger('fitz').setLevel(logging.ERROR)
+    logging.getLogger('ocrmypdf').setLevel(logging.ERROR)
+    warnings.simplefilter(action='ignore', category=FutureWarning)

marker/models.py

@@ -0,0 +1,58 @@
+from marker.postprocessors.editor import load_editing_model
+from surya.model.detection import segformer
+from texify.model.model import load_model as load_texify_model
+from texify.model.processor import load_processor as load_texify_processor
+from marker.settings import settings
+from surya.model.recognition.model import load_model as load_recognition_model
+from surya.model.recognition.processor import load_processor as load_recognition_processor
+from surya.model.ordering.model import load_model as load_order_model
+from surya.model.ordering.processor import load_processor as load_order_processor
+
+
+def setup_recognition_model(langs):
+    rec_model = load_recognition_model(langs=langs)
+    rec_processor = load_recognition_processor()
+    rec_model.processor = rec_processor
+    return rec_model
+
+
+def setup_detection_model():
+    model = segformer.load_model()
+    processor = segformer.load_processor()
+    model.processor = processor
+    return model
+
+
+def setup_texify_model():
+    texify_model = load_texify_model(checkpoint=settings.TEXIFY_MODEL_NAME, device=settings.TORCH_DEVICE_MODEL, dtype=settings.TEXIFY_DTYPE)
+    texify_processor = load_texify_processor()
+    texify_model.processor = texify_processor
+    return texify_model
+
+
+def setup_layout_model():
+    model = segformer.load_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
+    processor = segformer.load_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
+    model.processor = processor
+    return model
+
+
+def setup_order_model():
+    model = load_order_model()
+    processor = load_order_processor()
+    model.processor = processor
+    return model
+
+
+def load_all_models(langs=None):
+    # langs is optional list of languages to prune from recognition MoE model
+    detection = setup_detection_model()
+    layout = setup_layout_model()
+    order = setup_order_model()
+    edit = load_editing_model()
+
+    # Only load recognition model if we'll need it for all pdfs
+    ocr = setup_recognition_model(langs) if (settings.OCR_ENGINE == "surya" and settings.OCR_ALL_PAGES) else None
+    texify = setup_texify_model()
+    model_lst = [texify, layout, order, edit, detection, ocr]
+    return model_lst

marker/ocr/detection.py

@@ -0,0 +1,30 @@
+from typing import List
+
+from pypdfium2 import PdfDocument
+from surya.detection import batch_text_detection
+
+from marker.pdf.images import render_image
+from marker.schema.page import Page
+from marker.settings import settings
+
+
+def get_batch_size():
+    if settings.DETECTOR_BATCH_SIZE is not None:
+        return settings.DETECTOR_BATCH_SIZE
+    elif settings.TORCH_DEVICE_MODEL == "cuda":
+        return 4
+    return 4
+
+
+def surya_detection(doc: PdfDocument, pages: List[Page], det_model, batch_multiplier=1):
+    processor = det_model.processor
+    max_len = min(len(pages), len(doc))
+    images = [render_image(doc[pnum], dpi=settings.SURYA_DETECTOR_DPI) for pnum in range(max_len)]
+
+    predictions = batch_text_detection(images, det_model, processor, batch_size=get_batch_size() * batch_multiplier)
+    for (page, pred) in zip(pages, predictions):
+        page.text_lines = pred
+
+
+
+

marker/ocr/heuristics.py

@@ -0,0 +1,74 @@
+import re
+from typing import List
+
+from marker.ocr.utils import alphanum_ratio
+from marker.schema.bbox import rescale_bbox, box_intersection_pct
+from marker.schema.page import Page
+from marker.settings import settings
+
+
+def should_ocr_page(page: Page, no_text: bool):
+    detected_lines_found = detected_line_coverage(page)
+
+    # OCR page if we got minimal text, or if we got too many spaces
+    conditions = [
+        no_text , # Full doc has no text, and needs full OCR
+        (len(page.prelim_text) > 0 and detect_bad_ocr(page.prelim_text)),  # Bad OCR
+        detected_lines_found is False, # didn't extract text for all detected lines
+    ]
+
+    return any(conditions) or settings.OCR_ALL_PAGES
+
+
+def detect_bad_ocr(text, space_threshold=.7, newline_threshold=.6, alphanum_threshold=.3):
+    if len(text) == 0:
+        # Assume OCR failed if we have no text
+        return True
+
+    spaces = len(re.findall(r'\s+', text))
+    alpha_chars = len(re.sub(r'\s+', '', text))
+    if spaces / (alpha_chars + spaces) > space_threshold:
+        return True
+
+    newlines = len(re.findall(r'\n+', text))
+    non_newlines = len(re.sub(r'\n+', '', text))
+    if newlines / (newlines + non_newlines) > newline_threshold:
+        return True
+
+    if alphanum_ratio(text) < alphanum_threshold: # Garbled text
+        return True
+
+    invalid_chars = len([c for c in text if c in settings.INVALID_CHARS])
+    if invalid_chars > max(4.0, len(text) * .03):
+        return True
+
+    return False
+
+
+def no_text_found(pages: List[Page]):
+    full_text = ""
+    for page in pages:
+        full_text += page.prelim_text
+    return len(full_text.strip()) == 0
+
+
+def detected_line_coverage(page: Page, intersect_thresh=.5, detection_thresh=.65):
+    found_lines = 0
+    for detected_line in page.text_lines.bboxes:
+
+        # Get bbox and rescale to match dimensions of original page
+        detected_bbox = detected_line.bbox
+        detected_bbox = rescale_bbox(page.text_lines.image_bbox, page.bbox, detected_bbox)
+
+        total_intersection = 0
+        for block in page.blocks:
+            for line in block.lines:
+                intersection_pct = box_intersection_pct(detected_bbox, line.bbox)
+                total_intersection += intersection_pct
+        if total_intersection > intersect_thresh:
+            found_lines += 1
+
+    total_lines = len(page.text_lines.bboxes)
+    if total_lines == 0:
+        return False
+    return found_lines / total_lines > detection_thresh

marker/ocr/lang.py

@@ -0,0 +1,36 @@
+from typing import List
+
+from surya.languages import CODE_TO_LANGUAGE, LANGUAGE_TO_CODE
+from surya.model.recognition.tokenizer import _tokenize as lang_tokenize
+
+from marker.ocr.tesseract import LANGUAGE_TO_TESSERACT_CODE, TESSERACT_CODE_TO_LANGUAGE
+from marker.settings import settings
+
+
+def langs_to_ids(langs: List[str]):
+    unique_langs = list(set(langs))
+    _, lang_tokens = lang_tokenize("", unique_langs)
+    return lang_tokens
+
+
+def replace_langs_with_codes(langs):
+    if settings.OCR_ENGINE == "surya":
+        for i, lang in enumerate(langs):
+            if lang.title() in LANGUAGE_TO_CODE:
+                langs[i] = LANGUAGE_TO_CODE[lang.title()]
+    else:
+        for i, lang in enumerate(langs):
+            if lang in LANGUAGE_TO_CODE:
+                langs[i] = LANGUAGE_TO_TESSERACT_CODE[lang]
+    return langs
+
+
+def validate_langs(langs):
+    if settings.OCR_ENGINE == "surya":
+        for lang in langs:
+            if lang not in CODE_TO_LANGUAGE:
+                raise ValueError(f"Invalid language code {lang} for Surya OCR")
+    else:
+        for lang in langs:
+            if lang not in TESSERACT_CODE_TO_LANGUAGE:
+                raise ValueError(f"Invalid language code {lang} for Tesseract")

marker/ocr/recognition.py

@@ -0,0 +1,168 @@
+from itertools import repeat
+from typing import List, Optional, Dict
+
+import pypdfium2 as pdfium
+import io
+from concurrent.futures import ThreadPoolExecutor
+
+from surya.ocr import run_recognition
+
+from marker.models import setup_recognition_model
+from marker.ocr.heuristics import should_ocr_page, no_text_found, detect_bad_ocr
+from marker.ocr.lang import langs_to_ids
+from marker.pdf.images import render_image
+from marker.schema.page import Page
+from marker.schema.block import Block, Line, Span
+from marker.settings import settings
+from marker.pdf.extract_text import get_text_blocks
+
+
+def get_batch_size():
+    if settings.RECOGNITION_BATCH_SIZE is not None:
+        return settings.RECOGNITION_BATCH_SIZE
+    elif settings.TORCH_DEVICE_MODEL == "cuda":
+        return 32
+    elif settings.TORCH_DEVICE_MODEL == "mps":
+        return 32
+    return 32
+
+
+def run_ocr(doc, pages: List[Page], langs: List[str], rec_model, batch_multiplier=1) -> (List[Page], Dict):
+    ocr_pages = 0
+    ocr_success = 0
+    ocr_failed = 0
+    no_text = no_text_found(pages)
+    ocr_idxs = []
+    for pnum, page in enumerate(pages):
+        ocr_needed = should_ocr_page(page, no_text)
+        if ocr_needed:
+            ocr_idxs.append(pnum)
+            ocr_pages += 1
+
+    # No pages need OCR
+    if ocr_pages == 0:
+        return pages, {"ocr_pages": 0, "ocr_failed": 0, "ocr_success": 0, "ocr_engine": "none"}
+
+    ocr_method = settings.OCR_ENGINE
+    if ocr_method is None:
+        return pages, {"ocr_pages": 0, "ocr_failed": 0, "ocr_success": 0, "ocr_engine": "none"}
+    elif ocr_method == "surya":
+        # Load model just in time if we're not OCRing everything
+        del_rec_model = False
+        if rec_model is None:
+            lang_tokens = langs_to_ids(langs)
+            rec_model = setup_recognition_model(lang_tokens)
+            del_rec_model = True
+
+        new_pages = surya_recognition(doc, ocr_idxs, langs, rec_model, pages, batch_multiplier=batch_multiplier)
+
+        if del_rec_model:
+            del rec_model
+    elif ocr_method == "ocrmypdf":
+        new_pages = tesseract_recognition(doc, ocr_idxs, langs)
+    else:
+        raise ValueError(f"Unknown OCR method {ocr_method}")
+
+    for orig_idx, page in zip(ocr_idxs, new_pages):
+        if detect_bad_ocr(page.prelim_text) or len(page.prelim_text) == 0:
+            ocr_failed += 1
+        else:
+            ocr_success += 1
+            pages[orig_idx] = page
+
+    return pages, {"ocr_pages": ocr_pages, "ocr_failed": ocr_failed, "ocr_success": ocr_success, "ocr_engine": ocr_method}
+
+
+def surya_recognition(doc, page_idxs, langs: List[str], rec_model, pages: List[Page], batch_multiplier=1) -> List[Optional[Page]]:
+    images = [render_image(doc[pnum], dpi=settings.SURYA_OCR_DPI) for pnum in page_idxs]
+    processor = rec_model.processor
+    selected_pages = [p for i, p in enumerate(pages) if i in page_idxs]
+
+    surya_langs = [langs] * len(page_idxs)
+    detection_results = [p.text_lines.bboxes for p in selected_pages]
+    polygons = [[b.polygon for b in bboxes] for bboxes in detection_results]
+
+    results = run_recognition(images, surya_langs, rec_model, processor, polygons=polygons, batch_size=get_batch_size() * batch_multiplier)
+
+    new_pages = []
+    for (page_idx, result, old_page) in zip(page_idxs, results, selected_pages):
+        text_lines = old_page.text_lines
+        ocr_results = result.text_lines
+        blocks = []
+        for i, line in enumerate(ocr_results):
+            block = Block(
+                bbox=line.bbox,
+                pnum=page_idx,
+                lines=[Line(
+                    bbox=line.bbox,
+                    spans=[Span(
+                        text=line.text,
+                        bbox=line.bbox,
+                        span_id=f"{page_idx}_{i}",
+                        font="",
+                        font_weight=0,
+                        font_size=0,
+                    )
+                    ]
+                )]
+            )
+            blocks.append(block)
+        page = Page(
+            blocks=blocks,
+            pnum=page_idx,
+            bbox=result.image_bbox,
+            rotation=0,
+            text_lines=text_lines,
+            ocr_method="surya"
+        )
+        new_pages.append(page)
+    return new_pages
+
+
+def tesseract_recognition(doc, page_idxs, langs: List[str]) -> List[Optional[Page]]:
+    pdf_pages = generate_single_page_pdfs(doc, page_idxs)
+    with ThreadPoolExecutor(max_workers=settings.OCR_PARALLEL_WORKERS) as executor:
+        pages = list(executor.map(_tesseract_recognition, pdf_pages, repeat(langs, len(pdf_pages))))
+
+    return pages
+
+
+def generate_single_page_pdfs(doc, page_idxs) -> List[io.BytesIO]:
+    pdf_pages = []
+    for page_idx in page_idxs:
+        blank_doc = pdfium.PdfDocument.new()
+        blank_doc.import_pages(doc, pages=[page_idx])
+        assert len(blank_doc) == 1, "Failed to import page"
+
+        in_pdf = io.BytesIO()
+        blank_doc.save(in_pdf)
+        in_pdf.seek(0)
+        pdf_pages.append(in_pdf)
+    return pdf_pages
+
+
+def _tesseract_recognition(in_pdf, langs: List[str]) -> Optional[Page]:
+    import ocrmypdf
+    out_pdf = io.BytesIO()
+
+    ocrmypdf.ocr(
+        in_pdf,
+        out_pdf,
+        language=langs[0],
+        output_type="pdf",
+        redo_ocr=None,
+        force_ocr=True,
+        progress_bar=False,
+        optimize=False,
+        fast_web_view=1e6,
+        skip_big=15,  # skip images larger than 15 megapixels
+        tesseract_timeout=settings.TESSERACT_TIMEOUT,
+        tesseract_non_ocr_timeout=settings.TESSERACT_TIMEOUT,
+    )
+
+    new_doc = pdfium.PdfDocument(out_pdf.getvalue())
+
+    blocks, _ = get_text_blocks(new_doc, max_pages=1)
+    page = blocks[0]
+    page.ocr_method = "tesseract"
+    return page

marker/ocr/tesseract.py

@@ -0,0 +1,97 @@
+LANGUAGE_TO_TESSERACT_CODE = {
+    'Afrikaans': 'afr',
+    'Amharic': 'amh',
+    'Arabic': 'ara',
+    'Assamese': 'asm',
+    'Azerbaijani': 'aze',
+    'Belarusian': 'bel',
+    'Bulgarian': 'bul',
+    'Bengali': 'ben',
+    'Breton': 'bre',
+    'Bosnian': 'bos',
+    'Catalan': 'cat',
+    'Czech': 'ces',
+    'Welsh': 'cym',
+    'Danish': 'dan',
+    'German': 'deu',
+    'Greek': 'ell',
+    'English': 'eng',
+    'Esperanto': 'epo',
+    'Spanish': 'spa',
+    'Estonian': 'est',
+    'Basque': 'eus',
+    'Persian': 'fas',
+    'Finnish': 'fin',
+    'French': 'fra',
+    'Western Frisian': 'fry',
+    'Irish': 'gle',
+    'Scottish Gaelic': 'gla',
+    'Galician': 'glg',
+    'Gujarati': 'guj',
+    'Hausa': 'hau',
+    'Hebrew': 'heb',
+    'Hindi': 'hin',
+    'Croatian': 'hrv',
+    'Hungarian': 'hun',
+    'Armenian': 'hye',
+    'Indonesian': 'ind',
+    'Icelandic': 'isl',
+    'Italian': 'ita',
+    'Japanese': 'jpn',
+    'Javanese': 'jav',
+    'Georgian': 'kat',
+    'Kazakh': 'kaz',
+    'Khmer': 'khm',
+    'Kannada': 'kan',
+    'Korean': 'kor',
+    'Kurdish': 'kur',
+    'Kyrgyz': 'kir',
+    'Latin': 'lat',
+    'Lao': 'lao',
+    'Lithuanian': 'lit',
+    'Latvian': 'lav',
+    'Malagasy': 'mlg',
+    'Macedonian': 'mkd',
+    'Malayalam': 'mal',
+    'Mongolian': 'mon',
+    'Marathi': 'mar',
+    'Malay': 'msa',
+    'Burmese': 'mya',
+    'Nepali': 'nep',
+    'Dutch': 'nld',
+    'Norwegian': 'nor',
+    'Oromo': 'orm',
+    'Oriya': 'ori',
+    'Punjabi': 'pan',
+    'Polish': 'pol',
+    'Pashto': 'pus',
+    'Portuguese': 'por',
+    'Romanian': 'ron',
+    'Russian': 'rus',
+    'Sanskrit': 'san',
+    'Sindhi': 'snd',
+    'Sinhala': 'sin',
+    'Slovak': 'slk',
+    'Slovenian': 'slv',
+    'Somali': 'som',
+    'Albanian': 'sqi',
+    'Serbian': 'srp',
+    'Sundanese': 'sun',
+    'Swedish': 'swe',
+    'Swahili': 'swa',
+    'Tamil': 'tam',
+    'Telugu': 'tel',
+    'Thai': 'tha',
+    'Tagalog': 'tgl',
+    'Turkish': 'tur',
+    'Uyghur': 'uig',
+    'Ukrainian': 'ukr',
+    'Urdu': 'urd',
+    'Uzbek': 'uzb',
+    'Vietnamese': 'vie',
+    'Xhosa': 'xho',
+    'Yiddish': 'yid',
+    'Chinese': 'chi_sim',
+}
+
+TESSERACT_CODE_TO_LANGUAGE = {v:k for k,v in LANGUAGE_TO_TESSERACT_CODE.items()}

marker/ocr/utils.py

@@ -0,0 +1,10 @@
+def alphanum_ratio(text):
+    text = text.replace(" ", "")
+    text = text.replace("\n", "")
+    alphanumeric_count = sum([1 for c in text if c.isalnum()])
+
+    if len(text) == 0:
+        return 1
+
+    ratio = alphanumeric_count / len(text)
+    return ratio

marker/output.py

@@ -0,0 +1,39 @@
+import os
+import json
+
+
+def get_subfolder_path(out_folder, fname):
+    subfolder_name = fname.split(".")[0]
+    subfolder_path = os.path.join(out_folder, subfolder_name)
+    return subfolder_path
+
+
+def get_markdown_filepath(out_folder, fname):
+    subfolder_path = get_subfolder_path(out_folder, fname)
+    out_filename = fname.rsplit(".", 1)[0] + ".md"
+    out_filename = os.path.join(subfolder_path, out_filename)
+    return out_filename
+
+
+def markdown_exists(out_folder, fname):
+    out_filename = get_markdown_filepath(out_folder, fname)
+    return os.path.exists(out_filename)
+
+
+def save_markdown(out_folder, fname, full_text, images, out_metadata):
+    subfolder_path = get_subfolder_path(out_folder, fname)
+    os.makedirs(subfolder_path, exist_ok=True)
+
+    markdown_filepath = get_markdown_filepath(out_folder, fname)
+    out_meta_filepath = markdown_filepath.rsplit(".", 1)[0] + "_meta.json"
+
+    with open(markdown_filepath, "w+", encoding='utf-8') as f:
+        f.write(full_text)
+    with open(out_meta_filepath, "w+") as f:
+        f.write(json.dumps(out_metadata, indent=4))
+
+    for filename, image in images.items():
+        image_filepath = os.path.join(subfolder_path, filename)
+        image.save(image_filepath, "PNG")
+
+    return subfolder_path

marker/pdf/extract_text.py

@@ -0,0 +1,121 @@
+import os
+from typing import List, Optional, Dict
+
+import pypdfium2 as pdfium
+import pypdfium2.internal as pdfium_i
+
+from marker.pdf.utils import font_flags_decomposer
+from marker.settings import settings
+from marker.schema.block import Span, Line, Block
+from marker.schema.page import Page
+from pdftext.extraction import dictionary_output
+
+os.environ["TESSDATA_PREFIX"] = settings.TESSDATA_PREFIX
+
+
+def pdftext_format_to_blocks(page, pnum: int) -> Page:
+    page_blocks = []
+    span_id = 0
+    for block_idx, block in enumerate(page["blocks"]):
+        block_lines = []
+        for l in block["lines"]:
+            spans = []
+            for i, s in enumerate(l["spans"]):
+                block_text = s["text"]
+                # Remove trailing newlines and carriage returns (tesseract)
+                while len(block_text) > 0 and block_text[-1] in ["\n", "\r"]:
+                    block_text = block_text[:-1]
+
+                block_text = block_text.replace("-\n", "") # Remove hyphenated line breaks
+                span_obj = Span(
+                    text=block_text, # Remove end of line newlines, not spaces
+                    bbox=s["bbox"],
+                    span_id=f"{pnum}_{span_id}",
+                    font=f"{s['font']['name']}_{font_flags_decomposer(s['font']['flags'])}", # Add font flags to end of font
+                    font_weight=s["font"]["weight"],
+                    font_size=s["font"]["size"],
+                )
+                spans.append(span_obj)  # Text, bounding box, span id
+                span_id += 1
+            line_obj = Line(
+                spans=spans,
+                bbox=l["bbox"],
+            )
+            # Only select valid lines, with positive bboxes
+            if line_obj.area >= 0:
+                block_lines.append(line_obj)
+        block_obj = Block(
+            lines=block_lines,
+            bbox=block["bbox"],
+            pnum=pnum
+        )
+        # Only select blocks with lines
+        if len(block_lines) > 0:
+            page_blocks.append(block_obj)
+
+    page_bbox = page["bbox"]
+    page_width = abs(page_bbox[2] - page_bbox[0])
+    page_height = abs(page_bbox[3] - page_bbox[1])
+    rotation = page["rotation"]
+
+    # Flip width and height if rotated
+    if rotation == 90 or rotation == 270:
+        page_width, page_height = page_height, page_width
+
+    char_blocks = page["blocks"]
+    page_bbox = [0, 0, page_width, page_height]
+    out_page = Page(
+        blocks=page_blocks,
+        pnum=page["page"],
+        bbox=page_bbox,
+        rotation=rotation,
+        char_blocks=char_blocks
+    )
+    return out_page
+
+
+def get_text_blocks(doc, max_pages: Optional[int] = None) -> (List[Page], Dict):
+    toc = get_toc(doc)
+
+    page_range = range(len(doc))
+    if max_pages:
+        range_end = min(max_pages, len(doc))
+        page_range = range(range_end)
+
+    char_blocks = dictionary_output(doc, page_range=page_range, keep_chars=True)
+    marker_blocks = [pdftext_format_to_blocks(page, pnum) for pnum, page in enumerate(char_blocks)]
+
+    return marker_blocks, toc
+
+
+def naive_get_text(doc):
+    full_text = ""
+    for page_idx in range(len(doc)):
+        page = doc.get_page(page_idx)
+        text_page = page.get_textpage()
+        full_text += text_page.get_text_bounded() + "\n"
+    return full_text
+
+
+def get_toc(doc, max_depth=15):
+    toc = doc.get_toc(max_depth=max_depth)
+    toc_list = []
+    for item in toc:
+        list_item = {
+            "title": item.title,
+            "level": item.level,
+            "is_closed": item.is_closed,
+            "n_kids": item.n_kids,
+            "page_index": item.page_index,
+            "view_mode": pdfium_i.ViewmodeToStr.get(item.view_mode),
+            "view_pos": item.view_pos,
+        }
+        toc_list.append(list_item)
+    return toc_list
+
+
+def get_length_of_text(fname: str) -> int:
+    doc = pdfium.PdfDocument(fname)
+    text = naive_get_text(doc).strip()
+
+    return len(text)

marker/pdf/images.py

@@ -0,0 +1,27 @@
+import pypdfium2 as pdfium
+from pypdfium2 import PdfPage
+
+from marker.schema.page import Page
+from marker.schema.bbox import rescale_bbox
+from marker.settings import settings
+
+
+def render_image(page: pdfium.PdfPage, dpi):
+    image = page.render(
+        scale=dpi / 72,
+        draw_annots=False
+    ).to_pil()
+    image = image.convert("RGB")
+    return image
+
+
+def render_bbox_image(page_obj: PdfPage, page: Page, bbox):
+    png_image = render_image(page_obj, settings.IMAGE_DPI)
+    # Rescale original pdf bbox bounds to match png image size
+    png_bbox = [0, 0, png_image.size[0], png_image.size[1]]
+    rescaled_merged = rescale_bbox(page.bbox, png_bbox, bbox)
+
+    # Crop out only the equation image
+    png_image = png_image.crop(rescaled_merged)
+    png_image = png_image.convert("RGB")
+    return png_image

marker/pdf/utils.py

@@ -0,0 +1,75 @@
+from typing import Optional
+
+import filetype
+
+from marker.settings import settings
+
+
+def find_filetype(fpath):
+    kind = filetype.guess(fpath)
+    if kind is None:
+        print(f"Could not determine filetype for {fpath}")
+        return "other"
+
+    mimetype = kind.mime
+
+    # Get extensions from mimetype
+    # The mimetype is not always consistent, so use in to check the most common formats
+    if "pdf" in mimetype:
+        return "pdf"
+    elif mimetype in settings.SUPPORTED_FILETYPES:
+        return settings.SUPPORTED_FILETYPES[mimetype]
+    else:
+        print(f"Found nonstandard filetype {mimetype}")
+        return "other"
+
+
+def font_flags_decomposer(flags: Optional[int]) -> str:
+    if flags is None:
+        return ""
+
+    flag_descriptions = []
+    if flags & (1 << 0):  # PDFFONT_FIXEDPITCH
+        flag_descriptions.append("fixed_pitch")
+    if flags & (1 << 1):  # PDFFONT_SERIF
+        flag_descriptions.append("serif")
+    if flags & (1 << 2):  # PDFFONT_SYMBOLIC
+        flag_descriptions.append("symbolic")
+    if flags & (1 << 3):  # PDFFONT_SCRIPT
+        flag_descriptions.append("script")
+    if flags & (1 << 5):  # PDFFONT_NONSYMBOLIC
+        flag_descriptions.append("non_symbolic")
+    if flags & (1 << 6):  # PDFFONT_ITALIC
+        flag_descriptions.append("italic")
+    if flags & (1 << 16): # PDFFONT_ALLCAP
+        flag_descriptions.append("all_cap")
+    if flags & (1 << 17): # PDFFONT_SMALLCAP
+        flag_descriptions.append("small_cap")
+    if flags & (1 << 18): # PDFFONT_FORCEBOLD
+        flag_descriptions.append("bold")
+    if flags & (1 << 19): # PDFFONT_USEEXTERNATTR
+        flag_descriptions.append("use_extern_attr")
+
+    return "_".join(flag_descriptions)
+
+
+def sort_block_group(blocks, tolerance=1.25):
+    vertical_groups = {}
+    for block in blocks:
+        if hasattr(block, "bbox"):
+            bbox = block.bbox
+        else:
+            bbox = block["bbox"]
+
+        group_key = round(bbox[1] / tolerance) * tolerance
+        if group_key not in vertical_groups:
+            vertical_groups[group_key] = []
+        vertical_groups[group_key].append(block)
+
+    # Sort each group horizontally and flatten the groups into a single list
+    sorted_blocks = []
+    for _, group in sorted(vertical_groups.items()):
+        sorted_group = sorted(group, key=lambda x: x.bbox[0] if hasattr(x, "bbox") else x["bbox"][0])
+        sorted_blocks.extend(sorted_group)
+
+    return sorted_blocks

marker/postprocessors/editor.py

@@ -0,0 +1,116 @@
+from collections import defaultdict
+from itertools import chain
+from typing import Optional
+
+from marker.settings import settings
+import torch
+import torch.nn.functional as F
+from marker.postprocessors.t5 import T5ForTokenClassification, byt5_tokenize
+
+
+def get_batch_size():
+    if settings.EDITOR_BATCH_SIZE is not None:
+        return settings.EDITOR_BATCH_SIZE
+    elif settings.TORCH_DEVICE_MODEL == "cuda":
+        return 12
+    return 6
+
+
+def load_editing_model():
+    if not settings.ENABLE_EDITOR_MODEL:
+        return None
+
+    model = T5ForTokenClassification.from_pretrained(
+            settings.EDITOR_MODEL_NAME,
+            torch_dtype=settings.MODEL_DTYPE,
+        ).to(settings.TORCH_DEVICE_MODEL)
+    model.eval()
+
+    model.config.label2id = {
+        "equal": 0,
+        "delete": 1,
+        "newline-1": 2,
+        "space-1": 3,
+    }
+    model.config.id2label = {v: k for k, v in model.config.label2id.items()}
+    return model
+
+
+def edit_full_text(text: str, model: Optional[T5ForTokenClassification], batch_multiplier=1) -> (str, dict):
+    if not model:
+        return text, {}
+
+    batch_size = get_batch_size() * batch_multiplier
+    tokenized = byt5_tokenize(text, settings.EDITOR_MAX_LENGTH)
+    input_ids = tokenized["input_ids"]
+    char_token_lengths = tokenized["char_token_lengths"]
+
+    # Run model
+    token_masks = []
+    for i in range(0, len(input_ids), batch_size):
+        batch_input_ids = tokenized["input_ids"][i: i + batch_size]
+        batch_input_ids = torch.tensor(batch_input_ids, device=model.device)
+        batch_attention_mask = tokenized["attention_mask"][i: i + batch_size]
+        batch_attention_mask = torch.tensor(batch_attention_mask, device=model.device)
+        with torch.inference_mode():
+            predictions = model(batch_input_ids, attention_mask=batch_attention_mask)
+
+        logits = predictions.logits.cpu()
+
+        # If the max probability is less than a threshold, we assume it's a bad prediction
+        # We want to be conservative to not edit the text too much
+        probs = F.softmax(logits, dim=-1)
+        max_prob = torch.max(probs, dim=-1)
+        cutoff_prob = max_prob.values < settings.EDITOR_CUTOFF_THRESH
+        labels = logits.argmax(-1)
+        labels[cutoff_prob] = model.config.label2id["equal"]
+        labels = labels.squeeze().tolist()
+        if len(labels) == settings.EDITOR_MAX_LENGTH:
+            labels = [labels]
+        labels = list(chain.from_iterable(labels))
+        token_masks.extend(labels)
+
+    # List of characters in the text
+    flat_input_ids = list(chain.from_iterable(input_ids))
+
+    # Strip special tokens 0,1.  Keep unknown token, although it should never be used
+    assert len(token_masks) == len(flat_input_ids)
+    token_masks = [mask for mask, token in zip(token_masks, flat_input_ids) if token >= 2]
+
+    assert len(token_masks) == len(list(text.encode("utf-8")))
+
+    edit_stats = defaultdict(int)
+    out_text = []
+    start = 0
+    for i, char in enumerate(text):
+        char_token_length = char_token_lengths[i]
+        masks = token_masks[start: start + char_token_length]
+        labels = [model.config.id2label[mask] for mask in masks]
+        if all(l == "delete" for l in labels):
+            # If we delete whitespace, roll with it, otherwise ignore
+            if char.strip():
+                out_text.append(char)
+            else:
+                edit_stats["delete"] += 1
+        elif labels[0] == "newline-1":
+            out_text.append("\n")
+            out_text.append(char)
+            edit_stats["newline-1"] += 1
+        elif labels[0] == "space-1":
+            out_text.append(" ")
+            out_text.append(char)
+            edit_stats["space-1"] += 1
+        else:
+            out_text.append(char)
+            edit_stats["equal"] += 1
+
+        start += char_token_length
+
+    out_text = "".join(out_text)
+    return out_text, edit_stats
+
+
+
+
+
+

marker/postprocessors/markdown.py

@@ -0,0 +1,190 @@
+from marker.schema.merged import MergedLine, MergedBlock, FullyMergedBlock
+from marker.schema.page import Page
+import re
+import regex
+from typing import List
+
+
+def escape_markdown(text):
+    # List of characters that need to be escaped in markdown
+    characters_to_escape = r"[#]"
+    # Escape each of these characters with a backslash
+    escaped_text = re.sub(characters_to_escape, r'\\\g<0>', text)
+    return escaped_text
+
+
+def surround_text(s, char_to_insert):
+    leading_whitespace = re.match(r'^(\s*)', s).group(1)
+    trailing_whitespace = re.search(r'(\s*)$', s).group(1)
+    stripped_string = s.strip()
+    modified_string = char_to_insert + stripped_string + char_to_insert
+    final_string = leading_whitespace + modified_string + trailing_whitespace
+    return final_string
+
+
+def merge_spans(pages: List[Page]) -> List[List[MergedBlock]]:
+    merged_blocks = []
+    for page in pages:
+        page_blocks = []
+        for blocknum, block in enumerate(page.blocks):
+            block_lines = []
+            for linenum, line in enumerate(block.lines):
+                line_text = ""
+                if len(line.spans) == 0:
+                    continue
+                fonts = []
+                for i, span in enumerate(line.spans):
+                    font = span.font.lower()
+                    next_span = None
+                    next_idx = 1
+                    while len(line.spans) > i + next_idx:
+                        next_span = line.spans[i + next_idx]
+                        next_idx += 1
+                        if len(next_span.text.strip()) > 2:
+                            break
+
+                    fonts.append(font)
+                    span_text = span.text
+
+                    # Don't bold or italicize very short sequences
+                    # Avoid bolding first and last sequence so lines can be joined properly
+                    if len(span_text) > 3 and 0 < i < len(line.spans) - 1:
+                        if span.italic and (not next_span or not next_span.italic):
+                            span_text = surround_text(span_text, "*")
+                        elif span.bold and (not next_span or not next_span.bold):
+                            span_text = surround_text(span_text, "**")
+                    line_text += span_text
+                block_lines.append(MergedLine(
+                    text=line_text,
+                    fonts=fonts,
+                    bbox=line.bbox
+                ))
+            if len(block_lines) > 0:
+                page_blocks.append(MergedBlock(
+                    lines=block_lines,
+                    pnum=block.pnum,
+                    bbox=block.bbox,
+                    block_type=block.block_type
+                ))
+        merged_blocks.append(page_blocks)
+
+    return merged_blocks
+
+
+def block_surround(text, block_type):
+    if block_type == "Section-header":
+        if not text.startswith("#"):
+            text = "\n## " + text.strip().title() + "\n"
+    elif block_type == "Title":
+        if not text.startswith("#"):
+            text = "# " + text.strip().title() + "\n"
+    elif block_type == "Table":
+        text = "\n" + text + "\n"
+    elif block_type == "List-item":
+        text = escape_markdown(text)
+    elif block_type == "Code":
+        text = "\n```\n" + text + "\n```\n"
+    elif block_type == "Text":
+        text = escape_markdown(text)
+    elif block_type == "Formula":
+        if text.strip().startswith("$$") and text.strip().endswith("$$"):
+            text = text.strip()
+            text = "\n" + text + "\n"
+    return text
+
+
+def line_separator(line1, line2, block_type, is_continuation=False):
+    # Should cover latin-derived languages and russian
+    lowercase_letters = r'\p{Lo}|\p{Ll}|\d'
+    hyphens = r'-—¬'
+    # Remove hyphen in current line if next line and current line appear to be joined
+    hyphen_pattern = regex.compile(rf'.*[{lowercase_letters}][{hyphens}]\s?$', regex.DOTALL)
+    if line1 and hyphen_pattern.match(line1) and regex.match(rf"^\s?[{lowercase_letters}]", line2):
+        # Split on — or - from the right
+        line1 = regex.split(rf"[{hyphens}]\s?$", line1)[0]
+        return line1.rstrip() + line2.lstrip()
+
+    all_letters = r'\p{L}|\d'
+    sentence_continuations = r',;\(\—\"\'\*'
+    sentence_ends = r'。ๆ\.?!'
+    line_end_pattern = regex.compile(rf'.*[{lowercase_letters}][{sentence_continuations}]?\s?$', regex.DOTALL)
+    line_start_pattern = regex.compile(rf'^\s?[{all_letters}]', regex.DOTALL)
+    sentence_end_pattern = regex.compile(rf'.*[{sentence_ends}]\s?$', regex.DOTALL)
+
+    text_blocks = ["Text", "List-item", "Footnote", "Caption", "Figure"]
+    if block_type in ["Title", "Section-header"]:
+        return line1.rstrip() + " " + line2.lstrip()
+    elif block_type == "Formula":
+        return line1 + "\n" + line2
+    elif line_end_pattern.match(line1) and line_start_pattern.match(line2) and block_type in text_blocks:
+        return line1.rstrip() + " " + line2.lstrip()
+    elif is_continuation:
+        return line1.rstrip() + " " + line2.lstrip()
+    elif block_type in text_blocks and sentence_end_pattern.match(line1):
+        return line1 + "\n\n" + line2
+    elif block_type == "Table":
+        return line1 + "\n\n" + line2
+    else:
+        return line1 + "\n" + line2
+
+
+def block_separator(line1, line2, block_type1, block_type2):
+    sep = "\n"
+    if block_type1 == "Text":
+        sep = "\n\n"
+
+    return sep + line2
+
+
+def merge_lines(blocks: List[List[MergedBlock]]):
+    text_blocks = []
+    prev_type = None
+    prev_line = None
+    block_text = ""
+    block_type = ""
+
+    for page in blocks:
+        for block in page:
+            block_type = block.block_type
+            if block_type != prev_type and prev_type:
+                text_blocks.append(
+                    FullyMergedBlock(
+                        text=block_surround(block_text, prev_type),
+                        block_type=prev_type
+                    )
+                )
+                block_text = ""
+
+            prev_type = block_type
+            # Join lines in the block together properly
+            for i, line in enumerate(block.lines):
+                line_height = line.bbox[3] - line.bbox[1]
+                prev_line_height = prev_line.bbox[3] - prev_line.bbox[1] if prev_line else 0
+                prev_line_x = prev_line.bbox[0] if prev_line else 0
+                prev_line = line
+                is_continuation = line_height == prev_line_height and line.bbox[0] == prev_line_x
+                if block_text:
+                    block_text = line_separator(block_text, line.text, block_type, is_continuation)
+                else:
+                    block_text = line.text
+
+    # Append the final block
+    text_blocks.append(
+        FullyMergedBlock(
+            text=block_surround(block_text, prev_type),
+            block_type=block_type
+        )
+    )
+    return text_blocks
+
+
+def get_full_text(text_blocks):
+    full_text = ""
+    prev_block = None
+    for block in text_blocks:
+        if prev_block:
+            full_text += block_separator(prev_block.text, block.text, prev_block.block_type, block.block_type)
+        else:
+            full_text += block.text
+        prev_block = block
+    return full_text