commit b8bbed94ec63ca023914713f1cec90cc98059c28
Author: beep <beep@wimdupont.com>
Date: Sun, 5 Apr 2026 12:51:58 +0000
Add tiny LLM learning demos
Diffstat:
| A | .gitignore | | | 1 | + |
| A | LICENSE | | | 674 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | README.md | | | 264 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | corpus.txt | | | 25 | +++++++++++++++++++++++++ |
| A | tiny_lm.py | | | 306 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | tiny_modern_lm.py | | | 274 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | tiny_transformer_lm.py | | | 385 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
7 files changed, 1929 insertions(+), 0 deletions(-)
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1 @@
+__pycache__/
diff --git a/LICENSE b/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.
+
+ <one line to give the program's name and a brief idea of what it does.>
+ Copyright (C) 2024 Wim Dupont
+
+ 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:
+
+ <program> Copyright (C) 2024 Wim Dupont
+ 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>.
diff --git a/README.md b/README.md
@@ -0,0 +1,264 @@
+# tiny-llm-demo
+
+A very small language-model project meant to show the raw mechanics behind an
+LLM-like system without hiding them behind frameworks.
+
+These are not real LLMs. They are tiny plain-Python demos that progressively
+move from a simple neural next-token model toward transformer-like structure.
+The point is to make the moving parts visible:
+
+- text becomes tokens
+- tokens become vectors
+- vectors go through a small neural network
+- the network predicts the next token
+- training nudges weights so the prediction gets better
+
+## Files
+
+- `tiny_lm.py`: model, training loop, and text generation
+- `tiny_transformer_lm.py`: second version with a minimal self-attention block
+- `tiny_modern_lm.py`: forward-only demo of a more modern transformer block
+- `corpus.txt`: small training text
+
+## Run
+
+```bash
+cd /home/dude/repositories/beep/tiny-llm-demo
+python3 tiny_lm.py
+```
+
+Generate from a custom prompt:
+
+```bash
+python3 tiny_lm.py --prompt "language models " --sample-length 200
+```
+
+Run the attention-based version:
+
+```bash
+python3 tiny_transformer_lm.py
+```
+
+Show the learned attention weights over the prompt context:
+
+```bash
+python3 tiny_transformer_lm.py --show-attention
+```
+
+Inspect a more modern transformer-style block:
+
+```bash
+python3 tiny_modern_lm.py
+```
+
+## What the code is showing
+
+The model has:
+
+1. A character vocabulary
+2. Token embeddings
+3. A context window that keeps token positions separate
+4. A hidden layer with `tanh`
+5. An output layer that predicts the next character
+
+At each training step:
+
+1. Take a short slice of text as context
+2. Ask the model for the next-character probabilities
+3. Compare them to the real next character
+4. Compute the loss
+5. Backpropagate the error
+6. Update the weights with SGD
+
+This is the same high-level story as a real LLM, just drastically smaller and
+with a much simpler architecture.
+
+For the sake of a short demo, the script repeats the small corpus several times
+by default so the model can learn visible patterns faster.
+
+## What it is not showing
+
+- a full transformer stack
+- multi-head attention
+- parallel GPU training
+- distributed data loading
+- instruction tuning
+- RLHF / preference optimization
+- inference optimization
+
+Those are important in real systems, but this project is meant to answer the
+more basic question: what does it look like before all that complexity is added?
+
+## Why there are two versions
+
+`tiny_lm.py` is the simpler baseline.
+
+- It takes the whole context window and feeds it through a fixed learned network.
+- Every position matters only because the code gives it a slot in the input.
+- There is no dynamic decision about which earlier token to focus on.
+
+`tiny_transformer_lm.py` is the more transformer-like version.
+
+- Each position gets its own vector.
+- The last position forms a query.
+- Earlier positions produce keys and values.
+- Attention scores decide which earlier positions matter for the next prediction.
+
+That is the conceptual jump toward transformers: instead of using one fixed
+computation over the context, the model learns to route information dynamically
+based on token-to-token interactions.
+
+In this toy version you can inspect that directly with `--show-attention`,
+which prints the weight assigned to each character position in the current
+prompt context.
+
+`tiny_modern_lm.py` goes one step closer to a real LLM block.
+
+- layer norm before sublayers
+- causal multi-head self-attention
+- residual connection after attention
+- feed-forward MLP
+- residual connection after the MLP
+- output projection to next-token logits
+
+This is much closer to the shape of a modern transformer block, but it is
+forward-only. The earlier two scripts are the training demos.
+
+## Modern LLM Concepts
+
+If you want a compact mental model, these are the core ideas:
+
+1. Text is tokenized.
+2. Tokens become learned vectors called embeddings.
+3. Position information is added so order is preserved.
+4. The model applies many stacked transformer blocks.
+5. Each block uses self-attention so tokens can dynamically weigh other tokens.
+6. Each block also uses a feed-forward sublayer to transform each position.
+7. Residual connections help information and gradients flow through deep stacks.
+8. Layer normalization or RMSNorm helps stabilize training.
+9. The final vectors are projected into logits over the vocabulary.
+10. Training minimizes next-token prediction error over huge corpora.
+
+## What Each Script Represents
+
+`tiny_lm.py`
+
+- Shows raw neural next-token training.
+- The context is pushed through a fixed learned network.
+- Good for understanding random initialization, loss, backprop, and SGD.
+
+`tiny_transformer_lm.py`
+
+- Shows the key jump toward transformer behavior.
+- The model uses query/key/value attention to weight context positions dynamically.
+- Good for understanding why attention is different from a fixed MLP over context.
+
+`tiny_modern_lm.py`
+
+- Shows the architectural shape of a more modern LLM block.
+- Adds multi-head attention, layer norm, feed-forward layers, and residual paths.
+- Good for understanding how modern transformer blocks are composed.
+
+## What Is Still Missing From These Toys
+
+Even with the third script, these demos are still missing important parts of a
+real modern LLM:
+
+- many stacked transformer blocks instead of one tiny block
+- real training for the full modern block
+- subword tokenization such as BPE instead of character-level tokens
+- large-scale optimization such as AdamW, schedules, clipping, and mixed precision
+- large curated datasets and large-scale data pipelines
+- efficient inference features such as KV cache, batching, and quantization
+- post-training such as instruction tuning and preference optimization
+- deployment and serving systems
+
+So the rough ladder is:
+
+- `tiny_lm.py`: neural language model basics
+- `tiny_transformer_lm.py`: attention basics
+- `tiny_modern_lm.py`: modern transformer block shape
+
+That gets you much closer to a modern LLM conceptually, even though the scale
+and training sophistication are still far away.
+
+## Mapping To Real LLM Terms
+
+If you want to connect these demos to standard terminology, use this table:
+
+- `character vocabulary` / `stoi` / `itos`
+ Real term: tokenizer vocabulary
+ In real systems this is usually a subword vocabulary such as BPE tokens, not characters.
+
+- `token embeddings`
+ Real term: embedding layer
+ This is the learned lookup table that turns token IDs into dense vectors.
+
+- `positional embeddings`
+ Real term: positional encoding / positional embedding
+ This gives the model information about token order.
+
+- `next-character` or `next-token` prediction
+ Real term: autoregressive language modeling objective
+ This is the core pretraining task for a GPT-style model.
+
+- `output projection to logits`
+ Real term: LM head
+ This maps the final hidden state into scores over the vocabulary.
+
+- `fixed learned network over context` in `tiny_lm.py`
+ Real term: not a transformer block
+ This is just a baseline neural language model for learning the training loop.
+
+- `query`, `key`, and `value` in `tiny_transformer_lm.py`
+ Real term: self-attention mechanism
+ This is the core transformer idea that lets tokens weigh other tokens dynamically.
+
+- `multi-head attention` in `tiny_modern_lm.py`
+ Real term: multi-head self-attention
+ Real models use multiple heads so different relationships can be tracked in parallel.
+
+- `feed-forward MLP`
+ Real term: transformer feed-forward network, often called FFN or MLP block
+ This is the per-position transformation that sits alongside attention.
+
+- `layer norm`
+ Real term: LayerNorm or RMSNorm
+ Modern models rely on this for stable deep training.
+
+- `residual connection`
+ Real term: residual / skip connection
+ This helps preserve information and makes deep optimization work better.
+
+- `stacking many blocks`
+ Real term: transformer depth
+ A real LLM is many transformer blocks stacked on top of each other.
+
+- `training on corpus.txt`
+ Real term: pretraining corpus
+ A real LLM uses a vastly larger and more varied dataset.
+
+- `SGD in the toy scripts`
+ Real term: optimizer
+ Real models usually use AdamW or similar optimizers, not plain SGD.
+
+- `sample generation`
+ Real term: inference / decoding
+ Real systems add practical features such as KV cache, batching, and quantization.
+
+- `missing instruction tuning`
+ Real term: post-training / supervised fine-tuning
+ This is the stage where a pretrained model is taught to respond usefully to prompts.
+
+- `missing preference optimization`
+ Real term: RLHF, DPO, or related alignment methods
+ This is where models are shaped to better match human preferences and policy goals.
+
+The rough real-world sequence is:
+
+1. Build a tokenizer and vocabulary.
+2. Define a transformer architecture.
+3. Pretrain it on next-token prediction.
+4. Fine-tune or instruction-tune it for useful interaction.
+5. Apply preference and safety training.
+6. Optimize inference and deploy it.
diff --git a/corpus.txt b/corpus.txt
@@ -0,0 +1,25 @@
+language models predict the next token.
+language models start as code and random numbers.
+language models learn by repeated correction.
+
+the code defines the model.
+the weights start random.
+training changes the weights.
+prediction creates error.
+error updates the weights.
+
+humans write the tokenizer.
+humans write the training loop.
+humans choose the architecture.
+the model learns the numerical patterns.
+
+read tokens.
+predict next token.
+measure error.
+adjust weights.
+repeat.
+
+small models show the same loop in a simpler form.
+large models scale the same loop up.
+the structure is chosen by people.
+the behavior is shaped by training.
diff --git a/tiny_lm.py b/tiny_lm.py
@@ -0,0 +1,306 @@
+#!/usr/bin/env python3
+"""
+Tiny character-level neural language model in plain Python.
+
+This is intentionally small and readable. It is not a transformer and not a real
+LLM. It exists to show the basic mechanics of next-token training:
+
+- tokenize text
+- embed tokens
+- run a forward pass
+- compute softmax loss
+- backpropagate gradients
+- update weights
+- sample new text
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import random
+from dataclasses import dataclass
+from pathlib import Path
+
+
+def zeros(length: int) -> list[float]:
+ return [0.0 for _ in range(length)]
+
+
+def random_vector(length: int, scale: float, rng: random.Random) -> list[float]:
+ return [rng.uniform(-scale, scale) for _ in range(length)]
+
+
+def random_matrix(rows: int, cols: int, scale: float, rng: random.Random) -> list[list[float]]:
+ return [random_vector(cols, scale, rng) for _ in range(rows)]
+
+
+def softmax(logits: list[float]) -> list[float]:
+ peak = max(logits)
+ exps = [math.exp(value - peak) for value in logits]
+ total = sum(exps)
+ return [value / total for value in exps]
+
+
+@dataclass
+class Batch:
+ context_ids: list[int]
+ target_id: int
+
+
+class CharTokenizer:
+ def __init__(self, text: str) -> None:
+ chars = sorted(set(text))
+ self.stoi = {char: index for index, char in enumerate(chars)}
+ self.itos = {index: char for char, index in self.stoi.items()}
+
+ @property
+ def vocab_size(self) -> int:
+ return len(self.stoi)
+
+ def encode(self, text: str) -> list[int]:
+ return [self.stoi[char] for char in text]
+
+ def decode(self, token_ids: list[int]) -> str:
+ return "".join(self.itos[token_id] for token_id in token_ids)
+
+
+class TinyLanguageModel:
+ def __init__(
+ self,
+ vocab_size: int,
+ context_size: int,
+ embed_dim: int,
+ hidden_dim: int,
+ seed: int,
+ ) -> None:
+ rng = random.Random(seed)
+ self.vocab_size = vocab_size
+ self.context_size = context_size
+ self.embed_dim = embed_dim
+ self.hidden_dim = hidden_dim
+ self.input_dim = context_size * embed_dim
+
+ self.token_embed = random_matrix(vocab_size, embed_dim, 0.08, rng)
+ self.w1 = random_matrix(self.input_dim, hidden_dim, 0.08, rng)
+ self.b1 = zeros(hidden_dim)
+ self.w2 = random_matrix(hidden_dim, vocab_size, 0.08, rng)
+ self.b2 = zeros(vocab_size)
+
+ def forward(self, context_ids: list[int]) -> tuple[list[float], dict[str, list[float] | list[int]]]:
+ combined = zeros(self.input_dim)
+ for position, token_id in enumerate(context_ids):
+ token_vector = self.token_embed[token_id]
+ for dim in range(self.embed_dim):
+ combined[position * self.embed_dim + dim] = token_vector[dim]
+
+ hidden_pre = zeros(self.hidden_dim)
+ for hidden_idx in range(self.hidden_dim):
+ total = self.b1[hidden_idx]
+ for dim in range(self.input_dim):
+ total += combined[dim] * self.w1[dim][hidden_idx]
+ hidden_pre[hidden_idx] = total
+
+ hidden = [math.tanh(value) for value in hidden_pre]
+
+ logits = zeros(self.vocab_size)
+ for vocab_idx in range(self.vocab_size):
+ total = self.b2[vocab_idx]
+ for hidden_idx in range(self.hidden_dim):
+ total += hidden[hidden_idx] * self.w2[hidden_idx][vocab_idx]
+ logits[vocab_idx] = total
+
+ cache = {
+ "context_ids": context_ids,
+ "combined": combined,
+ "hidden": hidden,
+ }
+ return logits, cache
+
+ def train_step(self, batch: Batch, learning_rate: float) -> float:
+ logits, cache = self.forward(batch.context_ids)
+ probs = softmax(logits)
+ loss = -math.log(probs[batch.target_id] + 1e-12)
+
+ dlogits = probs[:]
+ dlogits[batch.target_id] -= 1.0
+
+ hidden = cache["hidden"]
+ combined = cache["combined"]
+
+ dw2 = [zeros(self.vocab_size) for _ in range(self.hidden_dim)]
+ db2 = dlogits[:]
+ for hidden_idx in range(self.hidden_dim):
+ for vocab_idx in range(self.vocab_size):
+ dw2[hidden_idx][vocab_idx] = hidden[hidden_idx] * dlogits[vocab_idx]
+
+ dhidden = zeros(self.hidden_dim)
+ for hidden_idx in range(self.hidden_dim):
+ total = 0.0
+ for vocab_idx in range(self.vocab_size):
+ total += self.w2[hidden_idx][vocab_idx] * dlogits[vocab_idx]
+ dhidden[hidden_idx] = total
+
+ dhidden_pre = zeros(self.hidden_dim)
+ for hidden_idx in range(self.hidden_dim):
+ dhidden_pre[hidden_idx] = dhidden[hidden_idx] * (1.0 - hidden[hidden_idx] ** 2)
+
+ dw1 = [zeros(self.hidden_dim) for _ in range(self.input_dim)]
+ db1 = dhidden_pre[:]
+ for dim in range(self.input_dim):
+ for hidden_idx in range(self.hidden_dim):
+ dw1[dim][hidden_idx] = combined[dim] * dhidden_pre[hidden_idx]
+
+ dcombined = zeros(self.input_dim)
+ for dim in range(self.input_dim):
+ total = 0.0
+ for hidden_idx in range(self.hidden_dim):
+ total += self.w1[dim][hidden_idx] * dhidden_pre[hidden_idx]
+ dcombined[dim] = total
+
+ for position, token_id in enumerate(batch.context_ids):
+ for dim in range(self.embed_dim):
+ gradient = dcombined[position * self.embed_dim + dim]
+ self.token_embed[token_id][dim] -= learning_rate * gradient
+
+ for dim in range(self.input_dim):
+ for hidden_idx in range(self.hidden_dim):
+ self.w1[dim][hidden_idx] -= learning_rate * dw1[dim][hidden_idx]
+ for hidden_idx in range(self.hidden_dim):
+ self.b1[hidden_idx] -= learning_rate * db1[hidden_idx]
+
+ for hidden_idx in range(self.hidden_dim):
+ for vocab_idx in range(self.vocab_size):
+ self.w2[hidden_idx][vocab_idx] -= learning_rate * dw2[hidden_idx][vocab_idx]
+ for vocab_idx in range(self.vocab_size):
+ self.b2[vocab_idx] -= learning_rate * db2[vocab_idx]
+
+ return loss
+
+ def sample_next_token(self, context_ids: list[int], rng: random.Random, temperature: float) -> int:
+ logits, _ = self.forward(context_ids)
+ scaled = [value / max(temperature, 1e-6) for value in logits]
+ probs = softmax(scaled)
+
+ threshold = rng.random()
+ running = 0.0
+ for index, prob in enumerate(probs):
+ running += prob
+ if threshold <= running:
+ return index
+ return len(probs) - 1
+
+
+def build_batches(token_ids: list[int], context_size: int) -> list[Batch]:
+ batches: list[Batch] = []
+ for index in range(len(token_ids) - context_size):
+ batches.append(
+ Batch(
+ context_ids=token_ids[index : index + context_size],
+ target_id=token_ids[index + context_size],
+ )
+ )
+ return batches
+
+
+def generate_text(
+ model: TinyLanguageModel,
+ tokenizer: CharTokenizer,
+ prompt: str,
+ sample_length: int,
+ rng: random.Random,
+ temperature: float,
+) -> str:
+ if not prompt:
+ prompt = "language "
+
+ if any(char not in tokenizer.stoi for char in prompt):
+ known = "".join(sorted(tokenizer.stoi.keys()))
+ raise ValueError(f"Prompt contains characters outside the training set. Known chars: {known!r}")
+
+ token_ids = tokenizer.encode(prompt)
+ if len(token_ids) < model.context_size:
+ token_ids = [token_ids[0]] * (model.context_size - len(token_ids)) + token_ids
+
+ generated = token_ids[:]
+ for _ in range(sample_length):
+ context_ids = generated[-model.context_size :]
+ next_token = model.sample_next_token(context_ids, rng, temperature)
+ generated.append(next_token)
+
+ return tokenizer.decode(generated)
+
+
+def load_corpus(path: Path) -> str:
+ return path.read_text(encoding="utf-8")
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description="Train a tiny next-character language model.")
+ parser.add_argument("--steps", type=int, default=3000, help="Number of SGD steps.")
+ parser.add_argument("--report-every", type=int, default=500, help="Progress report interval.")
+ parser.add_argument("--learning-rate", type=float, default=0.08, help="SGD learning rate.")
+ parser.add_argument("--context-size", type=int, default=12, help="Context window length.")
+ parser.add_argument("--embed-dim", type=int, default=24, help="Embedding size.")
+ parser.add_argument("--hidden-dim", type=int, default=48, help="Hidden layer size.")
+ parser.add_argument("--sample-length", type=int, default=220, help="Characters to generate.")
+ parser.add_argument("--temperature", type=float, default=0.9, help="Sampling temperature.")
+ parser.add_argument("--seed", type=int, default=7, help="Random seed.")
+ parser.add_argument("--prompt", default="language models ", help="Initial prompt for generation.")
+ parser.add_argument(
+ "--repeat-corpus",
+ type=int,
+ default=8,
+ help="Repeat the corpus this many times to make the tiny demo learn faster.",
+ )
+ args = parser.parse_args()
+
+ project_dir = Path(__file__).resolve().parent
+ corpus_text = load_corpus(project_dir / "corpus.txt")
+ corpus_text = corpus_text * max(args.repeat_corpus, 1)
+
+ tokenizer = CharTokenizer(corpus_text)
+ token_ids = tokenizer.encode(corpus_text)
+ batches = build_batches(token_ids, args.context_size)
+ if not batches:
+ raise ValueError("Corpus is too small for the requested context size.")
+
+ model = TinyLanguageModel(
+ vocab_size=tokenizer.vocab_size,
+ context_size=args.context_size,
+ embed_dim=args.embed_dim,
+ hidden_dim=args.hidden_dim,
+ seed=args.seed,
+ )
+
+ rng = random.Random(args.seed)
+ running_loss = 0.0
+
+ print(f"vocab size: {tokenizer.vocab_size}")
+ print(f"training samples: {len(batches)}")
+ print("training...")
+
+ for step in range(1, args.steps + 1):
+ batch = batches[rng.randrange(len(batches))]
+ loss = model.train_step(batch, args.learning_rate)
+ running_loss += loss
+
+ if step % args.report_every == 0 or step == 1:
+ avg_loss = running_loss / min(step, args.report_every)
+ print(f"step {step:5d} | avg loss {avg_loss:.4f}")
+ running_loss = 0.0
+
+ print("\n--- sample ---\n")
+ sample = generate_text(
+ model=model,
+ tokenizer=tokenizer,
+ prompt=args.prompt,
+ sample_length=args.sample_length,
+ rng=rng,
+ temperature=args.temperature,
+ )
+ print(sample)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tiny_modern_lm.py b/tiny_modern_lm.py
@@ -0,0 +1,274 @@
+#!/usr/bin/env python3
+"""
+Tiny forward-only demo of a more modern transformer-style language model block.
+
+This file is meant to show architecture, not serious training. It includes:
+
+- token embeddings
+- positional embeddings
+- layer normalization
+- causal multi-head self-attention
+- residual connections
+- feed-forward network
+- output projection to logits
+
+Why forward-only?
+Because implementing a faithful transformer block *and* all of its backpropagation
+in plain standard-library Python would swamp the core ideas in gradient code.
+The earlier demos show training mechanics. This file shows the modern block shape.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import random
+from pathlib import Path
+
+
+def zeros(length: int) -> list[float]:
+ return [0.0 for _ in range(length)]
+
+
+def random_vector(length: int, scale: float, rng: random.Random) -> list[float]:
+ return [rng.uniform(-scale, scale) for _ in range(length)]
+
+
+def random_matrix(rows: int, cols: int, scale: float, rng: random.Random) -> list[list[float]]:
+ return [random_vector(cols, scale, rng) for _ in range(rows)]
+
+
+def add(a: list[float], b: list[float]) -> list[float]:
+ return [x + y for x, y in zip(a, b)]
+
+
+def softmax(logits: list[float]) -> list[float]:
+ peak = max(logits)
+ exps = [math.exp(value - peak) for value in logits]
+ total = sum(exps)
+ return [value / total for value in exps]
+
+
+def gelu(values: list[float]) -> list[float]:
+ out = []
+ for x in values:
+ out.append(0.5 * x * (1.0 + math.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * x ** 3))))
+ return out
+
+
+def matvec(vector: list[float], matrix: list[list[float]]) -> list[float]:
+ cols = len(matrix[0])
+ out = zeros(cols)
+ for col in range(cols):
+ total = 0.0
+ for row in range(len(vector)):
+ total += vector[row] * matrix[row][col]
+ out[col] = total
+ return out
+
+
+def layer_norm(vector: list[float], gamma: list[float], beta: list[float], eps: float = 1e-5) -> list[float]:
+ mean = sum(vector) / len(vector)
+ variance = sum((value - mean) ** 2 for value in vector) / len(vector)
+ denom = math.sqrt(variance + eps)
+ return [((value - mean) / denom) * gamma[i] + beta[i] for i, value in enumerate(vector)]
+
+
+class CharTokenizer:
+ def __init__(self, text: str) -> None:
+ chars = sorted(set(text))
+ self.stoi = {char: index for index, char in enumerate(chars)}
+ self.itos = {index: char for char, index in self.stoi.items()}
+
+ @property
+ def vocab_size(self) -> int:
+ return len(self.stoi)
+
+ def encode(self, text: str) -> list[int]:
+ return [self.stoi[char] for char in text]
+
+ def decode(self, token_ids: list[int]) -> str:
+ return "".join(self.itos[token_id] for token_id in token_ids)
+
+
+class TinyModernTransformerLM:
+ def __init__(self, vocab_size: int, context_size: int, embed_dim: int, num_heads: int, seed: int) -> None:
+ if embed_dim % num_heads != 0:
+ raise ValueError("embed_dim must be divisible by num_heads")
+
+ rng = random.Random(seed)
+ self.vocab_size = vocab_size
+ self.context_size = context_size
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.head_dim = embed_dim // num_heads
+
+ self.token_embed = random_matrix(vocab_size, embed_dim, 0.08, rng)
+ self.pos_embed = random_matrix(context_size, embed_dim, 0.08, rng)
+
+ self.ln1_gamma = [1.0] * embed_dim
+ self.ln1_beta = [0.0] * embed_dim
+ self.ln2_gamma = [1.0] * embed_dim
+ self.ln2_beta = [0.0] * embed_dim
+
+ self.wq = random_matrix(embed_dim, embed_dim, 0.08, rng)
+ self.wk = random_matrix(embed_dim, embed_dim, 0.08, rng)
+ self.wv = random_matrix(embed_dim, embed_dim, 0.08, rng)
+ self.wo = random_matrix(embed_dim, embed_dim, 0.08, rng)
+
+ ff_dim = embed_dim * 4
+ self.w1 = random_matrix(embed_dim, ff_dim, 0.08, rng)
+ self.b1 = zeros(ff_dim)
+ self.w2 = random_matrix(ff_dim, embed_dim, 0.08, rng)
+ self.b2 = zeros(embed_dim)
+
+ self.lm_head = random_matrix(embed_dim, vocab_size, 0.08, rng)
+ self.lm_bias = zeros(vocab_size)
+
+ def causal_attention(self, x_vectors: list[list[float]]) -> tuple[list[list[float]], list[list[list[float]]]]:
+ q_all = [matvec(x, self.wq) for x in x_vectors]
+ k_all = [matvec(x, self.wk) for x in x_vectors]
+ v_all = [matvec(x, self.wv) for x in x_vectors]
+ scale = math.sqrt(self.head_dim)
+
+ head_weights: list[list[list[float]]] = []
+ head_outputs: list[list[list[float]]] = []
+
+ for head in range(self.num_heads):
+ start = head * self.head_dim
+ stop = start + self.head_dim
+ weights_for_head: list[list[float]] = []
+ outputs_for_head: list[list[float]] = []
+
+ for target_pos in range(self.context_size):
+ scores = []
+ for source_pos in range(self.context_size):
+ if source_pos > target_pos:
+ scores.append(-1e9)
+ continue
+ dot = 0.0
+ for dim in range(start, stop):
+ dot += q_all[target_pos][dim] * k_all[source_pos][dim]
+ scores.append(dot / scale)
+
+ attn = softmax(scores)
+ weights_for_head.append(attn)
+
+ out = zeros(self.head_dim)
+ for source_pos in range(self.context_size):
+ for local_dim, dim in enumerate(range(start, stop)):
+ out[local_dim] += attn[source_pos] * v_all[source_pos][dim]
+ outputs_for_head.append(out)
+
+ head_weights.append(weights_for_head)
+ head_outputs.append(outputs_for_head)
+
+ combined_outputs: list[list[float]] = []
+ for pos in range(self.context_size):
+ merged = []
+ for head in range(self.num_heads):
+ merged.extend(head_outputs[head][pos])
+ combined_outputs.append(matvec(merged, self.wo))
+
+ return combined_outputs, head_weights
+
+ def feed_forward(self, vector: list[float]) -> list[float]:
+ hidden = matvec(vector, self.w1)
+ hidden = [hidden[i] + self.b1[i] for i in range(len(hidden))]
+ hidden = gelu(hidden)
+ out = matvec(hidden, self.w2)
+ return [out[i] + self.b2[i] for i in range(len(out))]
+
+ def forward(self, context_ids: list[int]) -> tuple[list[float], dict[str, object]]:
+ x = []
+ for pos, token_id in enumerate(context_ids):
+ x.append(add(self.token_embed[token_id], self.pos_embed[pos]))
+
+ ln1 = [layer_norm(vec, self.ln1_gamma, self.ln1_beta) for vec in x]
+ attn_out, head_weights = self.causal_attention(ln1)
+ x = [add(x[pos], attn_out[pos]) for pos in range(self.context_size)]
+
+ ln2 = [layer_norm(vec, self.ln2_gamma, self.ln2_beta) for vec in x]
+ ff_out = [self.feed_forward(vec) for vec in ln2]
+ x = [add(x[pos], ff_out[pos]) for pos in range(self.context_size)]
+
+ final_state = x[-1]
+ logits = matvec(final_state, self.lm_head)
+ logits = [logits[i] + self.lm_bias[i] for i in range(self.vocab_size)]
+
+ cache = {
+ "head_weights": head_weights,
+ "final_state": final_state,
+ }
+ return logits, cache
+
+
+def top_k_indices(values: list[float], k: int) -> list[int]:
+ return sorted(range(len(values)), key=lambda index: values[index], reverse=True)[:k]
+
+
+def load_corpus(path: Path) -> str:
+ return path.read_text(encoding="utf-8")
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description="Inspect a tiny modern transformer-style LM block.")
+ parser.add_argument("--context-size", type=int, default=12, help="Context window length.")
+ parser.add_argument("--embed-dim", type=int, default=24, help="Embedding size.")
+ parser.add_argument("--num-heads", type=int, default=3, help="Number of attention heads.")
+ parser.add_argument("--seed", type=int, default=7, help="Random seed.")
+ parser.add_argument("--prompt", default="language mod", help="Prompt to inspect.")
+ parser.add_argument("--top-k", type=int, default=8, help="How many next-token candidates to print.")
+ args = parser.parse_args()
+
+ project_dir = Path(__file__).resolve().parent
+ corpus_text = load_corpus(project_dir / "corpus.txt")
+ tokenizer = CharTokenizer(corpus_text)
+
+ if any(char not in tokenizer.stoi for char in args.prompt):
+ known = "".join(sorted(tokenizer.stoi.keys()))
+ raise ValueError(f"Prompt contains characters outside the training set. Known chars: {known!r}")
+
+ token_ids = tokenizer.encode(args.prompt)
+ if len(token_ids) < args.context_size:
+ token_ids = [token_ids[0]] * (args.context_size - len(token_ids)) + token_ids
+ context_ids = token_ids[-args.context_size :]
+
+ model = TinyModernTransformerLM(
+ vocab_size=tokenizer.vocab_size,
+ context_size=args.context_size,
+ embed_dim=args.embed_dim,
+ num_heads=args.num_heads,
+ seed=args.seed,
+ )
+
+ logits, cache = model.forward(context_ids)
+ probs = softmax(logits)
+ top = top_k_indices(probs, args.top_k)
+
+ print("tiny modern transformer-style LM block")
+ print(f"context: {tokenizer.decode(context_ids)!r}")
+ print(f"vocab size: {tokenizer.vocab_size}")
+ print(f"embed dim: {args.embed_dim}")
+ print(f"heads: {args.num_heads}")
+
+ print("\n--- top next-token candidates ---\n")
+ for index in top:
+ char = tokenizer.itos[index]
+ label = "\\n" if char == "\n" else char
+ print(f"{label!r}: {probs[index]:.4f}")
+
+ head_weights = cache["head_weights"]
+ print("\n--- attention from final position ---\n")
+ context_text = tokenizer.decode(context_ids)
+ for head_index, weights_for_head in enumerate(head_weights):
+ final_weights = weights_for_head[-1]
+ print(f"head {head_index}:")
+ for pos, weight in enumerate(final_weights):
+ char = context_text[pos]
+ label = "\\n" if char == "\n" else char
+ print(f" pos {pos:2d} | char {label!r} | weight {weight:.4f}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tiny_transformer_lm.py b/tiny_transformer_lm.py
@@ -0,0 +1,385 @@
+#!/usr/bin/env python3
+"""
+Tiny character-level language model with a minimal causal self-attention block.
+
+This is still far from a real transformer, but it shows the key structural leap
+from the plain MLP demo:
+
+- each context position gets its own representation
+- the final position forms a query
+- all earlier positions provide keys and values
+- attention weights decide what matters dynamically
+
+The model predicts the next character from the attended summary of the context.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import random
+from dataclasses import dataclass
+from pathlib import Path
+
+
+def zeros(length: int) -> list[float]:
+ return [0.0 for _ in range(length)]
+
+
+def random_vector(length: int, scale: float, rng: random.Random) -> list[float]:
+ return [rng.uniform(-scale, scale) for _ in range(length)]
+
+
+def random_matrix(rows: int, cols: int, scale: float, rng: random.Random) -> list[list[float]]:
+ return [random_vector(cols, scale, rng) for _ in range(rows)]
+
+
+def softmax(logits: list[float]) -> list[float]:
+ peak = max(logits)
+ exps = [math.exp(value - peak) for value in logits]
+ total = sum(exps)
+ return [value / total for value in exps]
+
+
+@dataclass
+class Batch:
+ context_ids: list[int]
+ target_id: int
+
+
+class CharTokenizer:
+ def __init__(self, text: str) -> None:
+ chars = sorted(set(text))
+ self.stoi = {char: index for index, char in enumerate(chars)}
+ self.itos = {index: char for char, index in self.stoi.items()}
+
+ @property
+ def vocab_size(self) -> int:
+ return len(self.stoi)
+
+ def encode(self, text: str) -> list[int]:
+ return [self.stoi[char] for char in text]
+
+ def decode(self, token_ids: list[int]) -> str:
+ return "".join(self.itos[token_id] for token_id in token_ids)
+
+
+class TinyAttentionLanguageModel:
+ def __init__(self, vocab_size: int, context_size: int, embed_dim: int, seed: int) -> None:
+ rng = random.Random(seed)
+ self.vocab_size = vocab_size
+ self.context_size = context_size
+ self.embed_dim = embed_dim
+ self.scale = math.sqrt(embed_dim)
+
+ self.token_embed = random_matrix(vocab_size, embed_dim, 0.08, rng)
+ self.pos_embed = random_matrix(context_size, embed_dim, 0.08, rng)
+
+ self.wq = random_matrix(embed_dim, embed_dim, 0.08, rng)
+ self.wk = random_matrix(embed_dim, embed_dim, 0.08, rng)
+ self.wv = random_matrix(embed_dim, embed_dim, 0.08, rng)
+ self.wo = random_matrix(embed_dim, vocab_size, 0.08, rng)
+ self.bo = zeros(vocab_size)
+
+ def matvec(self, vector: list[float], matrix: list[list[float]]) -> list[float]:
+ cols = len(matrix[0])
+ out = zeros(cols)
+ for col in range(cols):
+ total = 0.0
+ for row in range(len(vector)):
+ total += vector[row] * matrix[row][col]
+ out[col] = total
+ return out
+
+ def forward(self, context_ids: list[int]) -> tuple[list[float], dict[str, object]]:
+ x_vectors: list[list[float]] = []
+ q_vectors: list[list[float]] = []
+ k_vectors: list[list[float]] = []
+ v_vectors: list[list[float]] = []
+
+ for position, token_id in enumerate(context_ids):
+ x = [
+ self.token_embed[token_id][dim] + self.pos_embed[position][dim]
+ for dim in range(self.embed_dim)
+ ]
+ x_vectors.append(x)
+ q_vectors.append(self.matvec(x, self.wq))
+ k_vectors.append(self.matvec(x, self.wk))
+ v_vectors.append(self.matvec(x, self.wv))
+
+ last_index = len(context_ids) - 1
+ q_last = q_vectors[last_index]
+
+ scores = zeros(len(context_ids))
+ for index in range(len(context_ids)):
+ dot = 0.0
+ for dim in range(self.embed_dim):
+ dot += q_last[dim] * k_vectors[index][dim]
+ scores[index] = dot / self.scale
+
+ attn = softmax(scores)
+
+ attended = zeros(self.embed_dim)
+ for index in range(len(context_ids)):
+ for dim in range(self.embed_dim):
+ attended[dim] += attn[index] * v_vectors[index][dim]
+
+ logits = zeros(self.vocab_size)
+ for vocab_idx in range(self.vocab_size):
+ total = self.bo[vocab_idx]
+ for dim in range(self.embed_dim):
+ total += attended[dim] * self.wo[dim][vocab_idx]
+ logits[vocab_idx] = total
+
+ cache: dict[str, object] = {
+ "context_ids": context_ids,
+ "x_vectors": x_vectors,
+ "q_vectors": q_vectors,
+ "k_vectors": k_vectors,
+ "v_vectors": v_vectors,
+ "scores": scores,
+ "attn": attn,
+ "attended": attended,
+ }
+ return logits, cache
+
+ def train_step(self, batch: Batch, learning_rate: float) -> float:
+ logits, cache = self.forward(batch.context_ids)
+ probs = softmax(logits)
+ loss = -math.log(probs[batch.target_id] + 1e-12)
+
+ dlogits = probs[:]
+ dlogits[batch.target_id] -= 1.0
+
+ attended = cache["attended"]
+ attn = cache["attn"]
+ x_vectors = cache["x_vectors"]
+ q_vectors = cache["q_vectors"]
+ k_vectors = cache["k_vectors"]
+ v_vectors = cache["v_vectors"]
+
+ dwo = [zeros(self.vocab_size) for _ in range(self.embed_dim)]
+ dbo = dlogits[:]
+ for dim in range(self.embed_dim):
+ for vocab_idx in range(self.vocab_size):
+ dwo[dim][vocab_idx] = attended[dim] * dlogits[vocab_idx]
+
+ dattended = zeros(self.embed_dim)
+ for dim in range(self.embed_dim):
+ total = 0.0
+ for vocab_idx in range(self.vocab_size):
+ total += self.wo[dim][vocab_idx] * dlogits[vocab_idx]
+ dattended[dim] = total
+
+ dattn = zeros(self.context_size)
+ dv_vectors = [zeros(self.embed_dim) for _ in range(self.context_size)]
+ for pos in range(self.context_size):
+ dot = 0.0
+ for dim in range(self.embed_dim):
+ dot += dattended[dim] * v_vectors[pos][dim]
+ dv_vectors[pos][dim] += attn[pos] * dattended[dim]
+ dattn[pos] = dot
+
+ weighted = 0.0
+ for pos in range(self.context_size):
+ weighted += attn[pos] * dattn[pos]
+
+ dscores = zeros(self.context_size)
+ for pos in range(self.context_size):
+ dscores[pos] = attn[pos] * (dattn[pos] - weighted)
+
+ dq_last = zeros(self.embed_dim)
+ dk_vectors = [zeros(self.embed_dim) for _ in range(self.context_size)]
+ for pos in range(self.context_size):
+ for dim in range(self.embed_dim):
+ dq_last[dim] += dscores[pos] * k_vectors[pos][dim] / self.scale
+ dk_vectors[pos][dim] += dscores[pos] * q_vectors[self.context_size - 1][dim] / self.scale
+
+ dwq = [zeros(self.embed_dim) for _ in range(self.embed_dim)]
+ dwk = [zeros(self.embed_dim) for _ in range(self.embed_dim)]
+ dwv = [zeros(self.embed_dim) for _ in range(self.embed_dim)]
+ dx_vectors = [zeros(self.embed_dim) for _ in range(self.context_size)]
+
+ last_x = x_vectors[self.context_size - 1]
+ for in_dim in range(self.embed_dim):
+ for out_dim in range(self.embed_dim):
+ dwq[in_dim][out_dim] += last_x[in_dim] * dq_last[out_dim]
+ dx_vectors[self.context_size - 1][in_dim] += self.wq[in_dim][out_dim] * dq_last[out_dim]
+
+ for pos in range(self.context_size):
+ x = x_vectors[pos]
+ for in_dim in range(self.embed_dim):
+ for out_dim in range(self.embed_dim):
+ dwk[in_dim][out_dim] += x[in_dim] * dk_vectors[pos][out_dim]
+ dwv[in_dim][out_dim] += x[in_dim] * dv_vectors[pos][out_dim]
+ dx_vectors[pos][in_dim] += self.wk[in_dim][out_dim] * dk_vectors[pos][out_dim]
+ dx_vectors[pos][in_dim] += self.wv[in_dim][out_dim] * dv_vectors[pos][out_dim]
+
+ for pos, token_id in enumerate(batch.context_ids):
+ for dim in range(self.embed_dim):
+ gradient = dx_vectors[pos][dim]
+ self.token_embed[token_id][dim] -= learning_rate * gradient
+ self.pos_embed[pos][dim] -= learning_rate * gradient
+
+ for in_dim in range(self.embed_dim):
+ for out_dim in range(self.embed_dim):
+ self.wq[in_dim][out_dim] -= learning_rate * dwq[in_dim][out_dim]
+ self.wk[in_dim][out_dim] -= learning_rate * dwk[in_dim][out_dim]
+ self.wv[in_dim][out_dim] -= learning_rate * dwv[in_dim][out_dim]
+
+ for dim in range(self.embed_dim):
+ for vocab_idx in range(self.vocab_size):
+ self.wo[dim][vocab_idx] -= learning_rate * dwo[dim][vocab_idx]
+ for vocab_idx in range(self.vocab_size):
+ self.bo[vocab_idx] -= learning_rate * dbo[vocab_idx]
+
+ return loss
+
+ def sample_next_token(self, context_ids: list[int], rng: random.Random, temperature: float) -> int:
+ logits, _ = self.forward(context_ids)
+ scaled = [value / max(temperature, 1e-6) for value in logits]
+ probs = softmax(scaled)
+
+ threshold = rng.random()
+ running = 0.0
+ for index, prob in enumerate(probs):
+ running += prob
+ if threshold <= running:
+ return index
+ return len(probs) - 1
+
+ def attention_weights(self, context_ids: list[int]) -> list[float]:
+ _, cache = self.forward(context_ids)
+ return list(cache["attn"])
+
+
+def build_batches(token_ids: list[int], context_size: int) -> list[Batch]:
+ batches: list[Batch] = []
+ for index in range(len(token_ids) - context_size):
+ batches.append(
+ Batch(
+ context_ids=token_ids[index : index + context_size],
+ target_id=token_ids[index + context_size],
+ )
+ )
+ return batches
+
+
+def generate_text(
+ model: TinyAttentionLanguageModel,
+ tokenizer: CharTokenizer,
+ prompt: str,
+ sample_length: int,
+ rng: random.Random,
+ temperature: float,
+) -> str:
+ if not prompt:
+ prompt = "language "
+
+ if any(char not in tokenizer.stoi for char in prompt):
+ known = "".join(sorted(tokenizer.stoi.keys()))
+ raise ValueError(f"Prompt contains characters outside the training set. Known chars: {known!r}")
+
+ token_ids = tokenizer.encode(prompt)
+ if len(token_ids) < model.context_size:
+ token_ids = [token_ids[0]] * (model.context_size - len(token_ids)) + token_ids
+
+ generated = token_ids[:]
+ for _ in range(sample_length):
+ context_ids = generated[-model.context_size :]
+ next_token = model.sample_next_token(context_ids, rng, temperature)
+ generated.append(next_token)
+
+ return tokenizer.decode(generated)
+
+
+def load_corpus(path: Path) -> str:
+ return path.read_text(encoding="utf-8")
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description="Train a tiny self-attention language model.")
+ parser.add_argument("--steps", type=int, default=2500, help="Number of SGD steps.")
+ parser.add_argument("--report-every", type=int, default=500, help="Progress report interval.")
+ parser.add_argument("--learning-rate", type=float, default=0.03, help="SGD learning rate.")
+ parser.add_argument("--context-size", type=int, default=8, help="Context window length.")
+ parser.add_argument("--embed-dim", type=int, default=16, help="Embedding size.")
+ parser.add_argument("--sample-length", type=int, default=220, help="Characters to generate.")
+ parser.add_argument("--temperature", type=float, default=0.8, help="Sampling temperature.")
+ parser.add_argument("--seed", type=int, default=7, help="Random seed.")
+ parser.add_argument("--prompt", default="language models ", help="Initial prompt for generation.")
+ parser.add_argument(
+ "--repeat-corpus",
+ type=int,
+ default=12,
+ help="Repeat the corpus this many times to make the tiny demo learn faster.",
+ )
+ parser.add_argument(
+ "--show-attention",
+ action="store_true",
+ help="Print attention weights for the final prompt context after training.",
+ )
+ args = parser.parse_args()
+
+ project_dir = Path(__file__).resolve().parent
+ corpus_text = load_corpus(project_dir / "corpus.txt")
+ corpus_text = corpus_text * max(args.repeat_corpus, 1)
+
+ tokenizer = CharTokenizer(corpus_text)
+ token_ids = tokenizer.encode(corpus_text)
+ batches = build_batches(token_ids, args.context_size)
+ if not batches:
+ raise ValueError("Corpus is too small for the requested context size.")
+
+ model = TinyAttentionLanguageModel(
+ vocab_size=tokenizer.vocab_size,
+ context_size=args.context_size,
+ embed_dim=args.embed_dim,
+ seed=args.seed,
+ )
+
+ rng = random.Random(args.seed)
+ running_loss = 0.0
+
+ print(f"vocab size: {tokenizer.vocab_size}")
+ print(f"training samples: {len(batches)}")
+ print("training...")
+
+ for step in range(1, args.steps + 1):
+ batch = batches[rng.randrange(len(batches))]
+ loss = model.train_step(batch, args.learning_rate)
+ running_loss += loss
+
+ if step % args.report_every == 0 or step == 1:
+ avg_loss = running_loss / min(step, args.report_every)
+ print(f"step {step:5d} | avg loss {avg_loss:.4f}")
+ running_loss = 0.0
+
+ print("\n--- sample ---\n")
+ sample = generate_text(
+ model=model,
+ tokenizer=tokenizer,
+ prompt=args.prompt,
+ sample_length=args.sample_length,
+ rng=rng,
+ temperature=args.temperature,
+ )
+ print(sample)
+
+ if args.show_attention:
+ prompt_ids = tokenizer.encode(args.prompt)
+ if len(prompt_ids) < model.context_size:
+ prompt_ids = [prompt_ids[0]] * (model.context_size - len(prompt_ids)) + prompt_ids
+ context_ids = prompt_ids[-model.context_size :]
+ weights = model.attention_weights(context_ids)
+
+ print("\n--- attention on final prompt context ---\n")
+ context_text = tokenizer.decode(context_ids)
+ for index, weight in enumerate(weights):
+ char = context_text[index]
+ label = "\\n" if char == "\n" else char
+ print(f"pos {index:2d} | char {label!r} | weight {weight:.4f}")
+
+
+if __name__ == "__main__":
+ main()
