Initial commit

README.md

@@ -0,0 +1,475 @@
+---
+language:
+- en
+license:
+- apache-2.0
+- bsd-3-clause
+tags:
+- summarization
+- led
+- summary
+- longformer
+- booksum
+- long-document
+- long-form
+datasets:
+- kmfoda/booksum
+metrics:
+- rouge
+widget:
+- text: large earthquakes along a given fault segment do not occur at random intervals
+    because it takes time to accumulate the strain energy for the rupture. The rates
+    at which tectonic plates move and accumulate strain at their boundaries are approximately
+    uniform. Therefore, in first approximation, one may expect that large ruptures
+    of the same fault segment will occur at approximately constant time intervals.
+    If subsequent main shocks have different amounts of slip across the fault, then
+    the recurrence time may vary, and the basic idea of periodic mainshocks must be
+    modified. For great plate boundary ruptures the length and slip often vary by
+    a factor of 2. Along the southern segment of the San Andreas fault the recurrence
+    interval is 145 years with variations of several decades. The smaller the standard
+    deviation of the average recurrence interval, the more specific could be the long
+    term prediction of a future mainshock.
+  example_title: earthquakes
+- text: ' A typical feed-forward neural field algorithm. Spatiotemporal coordinates
+    are fed into a neural network that predicts values in the reconstructed domain.
+    Then, this domain is mapped to the sensor domain where sensor measurements are
+    available as supervision. Class and Section Problems Addressed Generalization
+    (Section 2) Inverse problems, ill-posed problems, editability; symmetries. Hybrid
+    Representations (Section 3) Computation & memory efficiency, representation capacity,
+    editability: Forward Maps (Section 4) Inverse problems Network Architecture (Section
+    5) Spectral bias, integration & derivatives. Manipulating Neural Fields (Section
+    6) Edit ability, constraints, regularization. Table 2: The five classes of techniques
+    in the neural field toolbox each addresses problems that arise in learning, inference,
+    and control. (Section 3). We can supervise reconstruction via differentiable forward
+    maps that transform Or project our domain (e.g, 3D reconstruction via 2D images;
+    Section 4) With appropriate network architecture choices, we can overcome neural
+    network spectral biases (blurriness) and efficiently compute derivatives and integrals
+    (Section 5). Finally, we can manipulate neural fields to add constraints and regularizations,
+    and to achieve editable representations (Section 6). Collectively, these classes
+    constitute a ''toolbox'' of techniques to help solve problems with neural fields
+    There are three components in a conditional neural field: (1) An encoder or inference
+    function € that outputs the conditioning latent variable 2 given an observation
+    0 E(0) =2. 2 is typically a low-dimensional vector, and is often referred to aS
+    a latent code Or feature code_ (2) A mapping function 4 between Z and neural field
+    parameters O: Y(z) = O; (3) The neural field itself $. The encoder € finds the
+    most probable z given the observations O: argmaxz P(2/0). The decoder maximizes
+    the inverse conditional probability to find the most probable 0 given Z: arg-
+    max P(Olz). We discuss different encoding schemes with different optimality guarantees
+    (Section 2.1.1), both global and local conditioning (Section 2.1.2), and different
+    mapping functions Y (Section 2.1.3) 2. Generalization Suppose we wish to estimate
+    a plausible 3D surface shape given a partial or noisy point cloud. We need a suitable
+    prior over the sur- face in its reconstruction domain to generalize to the partial
+    observations. A neural network expresses a prior via the function space of its
+    architecture and parameters 0, and generalization is influenced by the inductive
+    bias of this function space (Section 5).'
+  example_title: scientific paper
+- text: ' the big variety of data coming from diverse sources is one of the key properties
+    of the big data phenomenon. It is, therefore, beneficial to understand how data
+    is generated in various environments and scenarios, before looking at what should
+    be done with this data and how to design the best possible architecture to accomplish
+    this The evolution of IT architectures, described in Chapter 2, means that the
+    data is no longer processed by a few big monolith systems, but rather by a group
+    of services In parallel to the processing layer, the underlying data storage has
+    also changed and became more distributed This, in turn, required a significant
+    paradigm shift as the traditional approach to transactions (ACID) could no longer
+    be supported. On top of this, cloud computing is becoming a major approach with
+    the benefits of reducing costs and providing on-demand scalability but at the
+    same time introducing concerns about privacy, data ownership, etc In the meantime
+    the Internet continues its exponential growth: Every day both structured and unstructured
+    data is published and available for processing: To achieve competitive advantage
+    companies have to relate their corporate resources to external services, e.g.
+    financial markets, weather forecasts, social media, etc While several of the sites
+    provide some sort of API to access the data in a more orderly fashion; countless
+    sources require advanced web mining and Natural Language Processing (NLP) processing
+    techniques: Advances in science push researchers to construct new instruments
+    for observing the universe O conducting experiments to understand even better
+    the laws of physics and other domains. Every year humans have at their disposal
+    new telescopes, space probes, particle accelerators, etc These instruments generate
+    huge streams of data, which need to be stored and analyzed. The constant drive
+    for efficiency in the industry motivates the introduction of new automation techniques
+    and process optimization: This could not be done without analyzing the precise
+    data that describe these processes. As more and more human tasks are automated,
+    machines provide rich data sets, which can be analyzed in real-time to drive efficiency
+    to new levels. Finally, it is now evident that the growth of the Internet of Things
+    is becoming a major source of data. More and more of the devices are equipped
+    with significant computational power and can generate a continuous data stream
+    from their sensors. In the subsequent sections of this chapter, we will look at
+    the domains described above to see what they generate in terms of data sets. We
+    will compare the volumes but will also look at what is characteristic and important
+    from their respective points of view. 3.1 The Internet is undoubtedly the largest
+    database ever created by humans. While several well described; cleaned, and structured
+    data sets have been made available through this medium, most of the resources
+    are of an ambiguous, unstructured, incomplete or even erroneous nature. Still,
+    several examples in the areas such as opinion mining, social media analysis, e-governance,
+    etc, clearly show the potential lying in these resources. Those who can successfully
+    mine and interpret the Internet data can gain unique insight and competitive advantage
+    in their business An important area of data analytics on the edge of corporate
+    IT and the Internet is Web Analytics.'
+  example_title: data science textbook
+- text: 'Transformer-based models have shown to be very useful for many NLP tasks.
+    However, a major limitation of transformers-based models is its O(n^2)O(n 2) time
+    & memory complexity (where nn is sequence length). Hence, it''s computationally
+    very expensive to apply transformer-based models on long sequences n > 512n>512.
+    Several recent papers, e.g. Longformer, Performer, Reformer, Clustered attention
+    try to remedy this problem by approximating the full attention matrix. You can
+    checkout 🤗''s recent blog post in case you are unfamiliar with these models.
+
+    BigBird (introduced in paper) is one of such recent models to address this issue.
+    BigBird relies on block sparse attention instead of normal attention (i.e. BERT''s
+    attention) and can handle sequences up to a length of 4096 at a much lower computational
+    cost compared to BERT. It has achieved SOTA on various tasks involving very long
+    sequences such as long documents summarization, question-answering with long contexts.
+
+    BigBird RoBERTa-like model is now available in 🤗Transformers. The goal of this
+    post is to give the reader an in-depth understanding of big bird implementation
+    & ease one''s life in using BigBird with 🤗Transformers. But, before going into
+    more depth, it is important to remember that the BigBird''s attention is an approximation
+    of BERT''s full attention and therefore does not strive to be better than BERT''s
+    full attention, but rather to be more efficient. It simply allows to apply transformer-based
+    models to much longer sequences since BERT''s quadratic memory requirement quickly
+    becomes unbearable. Simply put, if we would have ∞ compute & ∞ time, BERT''s attention
+    would be preferred over block sparse attention (which we are going to discuss
+    in this post).
+
+    If you wonder why we need more compute when working with longer sequences, this
+    blog post is just right for you!
+
+    Some of the main questions one might have when working with standard BERT-like
+    attention include:
+
+    Do all tokens really have to attend to all other tokens? Why not compute attention
+    only over important tokens? How to decide what tokens are important? How to attend
+    to just a few tokens in a very efficient way? In this blog post, we will try to
+    answer those questions.
+
+    What tokens should be attended to? We will give a practical example of how attention
+    works by considering the sentence ''BigBird is now available in HuggingFace for
+    extractive question answering''. In BERT-like attention, every word would simply
+    attend to all other tokens.
+
+    Let''s think about a sensible choice of key tokens that a queried token actually
+    only should attend to by writing some pseudo-code. Will will assume that the token
+    available is queried and build a sensible list of key tokens to attend to.
+
+    >>> # let''s consider following sentence as an example >>> example = [''BigBird'',
+    ''is'', ''now'', ''available'', ''in'', ''HuggingFace'', ''for'', ''extractive'',
+    ''question'', ''answering'']
+
+    >>> # further let''s assume, we''re trying to understand the representation of
+    ''available'' i.e. >>> query_token = ''available'' >>> # We will initialize an
+    empty `set` and fill up the tokens of our interest as we proceed in this section.
+    >>> key_tokens = [] # => currently ''available'' token doesn''t have anything
+    to attend Nearby tokens should be important because, in a sentence (sequence of
+    words), the current word is highly dependent on neighboring past & future tokens.
+    This intuition is the idea behind the concept of sliding attention.'
+  example_title: bigbird blog intro
+- text: 'The majority of available text summarization datasets include short-form
+    source documents that lack long-range causal and temporal dependencies, and often
+    contain strong layout and stylistic biases. While relevant, such datasets will
+    offer limited challenges for future generations of text summarization systems.
+    We address these issues by introducing BookSum, a collection of datasets for long-form
+    narrative summarization. Our dataset covers source documents from the literature
+    domain, such as novels, plays and stories, and includes highly abstractive, human
+    written summaries on three levels of granularity of increasing difficulty: paragraph-,
+    chapter-, and book-level. The domain and structure of our dataset poses a unique
+    set of challenges for summarization systems, which include: processing very long
+    documents, non-trivial causal and temporal dependencies, and rich discourse structures.
+    To facilitate future work, we trained and evaluated multiple extractive and abstractive
+    summarization models as baselines for our dataset.'
+  example_title: BookSum Abstract
+inference:
+  parameters:
+    max_length: 64
+    min_length: 8
+    no_repeat_ngram_size: 3
+    early_stopping: true
+    repetition_penalty: 3.5
+    length_penalty: 0.3
+    encoder_no_repeat_ngram_size: 3
+    num_beams: 4
+model-index:
+- name: pszemraj/led-large-book-summary
+  results:
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: kmfoda/booksum
+      type: kmfoda/booksum
+      config: kmfoda--booksum
+      split: test
+    metrics:
+    - type: rouge
+      value: 31.7308
+      name: ROUGE-1
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiNjJmZjMxYTY0OGU3MzNjNmIzNmYyODNlNDg2ZGRhZDAzNTMwMDM5YWMxODc1OTc1ZWE3MzM2OTg1ODFhZDBkNCIsInZlcnNpb24iOjF9.B8BCKgySYVZW910_1zP0LfCpQYJbAe6loyWut76JlgZb2kV1_x9ybqtNESX0ka-lNqhYyXUNDpuS-7pTmsJVDg
+    - type: rouge
+      value: 5.3311
+      name: ROUGE-2
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiYzViMmY4ODFjYTc5ODk5MmRhMDQ3ZDRiYWQwMDg0OTk3ZTA4NDAxYTNiNDgyMmI4NDA3ZDMwYWViOTBkODBjNyIsInZlcnNpb24iOjF9.MOhJLDcgvv93mVFL1igIgIiTAH3b2Xa4gmBObq7RF44Mmu8Kxtd1KP7rOlDVFOrtrsooGPGsyE1GMCQ2kqeMDg
+    - type: rouge
+      value: 16.1465
+      name: ROUGE-L
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiNzNjMzEwMTliZGE3ZmQ4M2UxMDAyMTY3YzJjZmMyMDYyN2YyNDM0N2VhNzI1MDc1YTg4MTRjMmEzNjVkNTk1NCIsInZlcnNpb24iOjF9.XLJ-DVKiYLlbw5E5rWADKbzUzf5fNHhlTCWPCC5dU4NI9Yeh76aR7TPt36ZzLDwTBknnR8KHqlaF8F8YAvBUAg
+    - type: rouge
+      value: 29.0883
+      name: ROUGE-LSUM
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiMTcwNzEwMmE5NjQxZTkzYmQyZDZmNzllYzYyNGI5OTMyNWMwNjdiM2I2YmM5YjdmY2E5OWQ3OTk3ZDA1MTc3YyIsInZlcnNpb24iOjF9.d6rFxjCB6RJNI_pn2DNNSjuZe4rdvj0RatkaTJRp5lP0F_AFfU5Zn9zRWzZJV7V-xMauIc4UhfdoLp9r_-CABA
+    - type: loss
+      value: 4.815707206726074
+      name: loss
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiNTMwMTgxMmJkODY3MjkzOWJhMzJhOTIxMWVkODhjZmM0MWUzMWQ1N2JkZjRhOTQxNmU1YWVjYzQ0MDNlZWI3OSIsInZlcnNpb24iOjF9.mkBQHYhYFfDV6F4klXGJ1dSsF-pbCs-6F9zcw6IYznwmXUjtk7m5J4Zt4JAju5LKz4YizvEcUCl_L0WddnfvDA
+    - type: gen_len
+      value: 154.9036
+      name: gen_len
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiMTc0ZmM1ZDM4MDE0MzY3MDM3OWJhNDkzZjJkZDdkMjU5M2JmMDJjYTIxODA1OTllNmY5ZWQzZDlmNWFiYzk4NiIsInZlcnNpb24iOjF9.VQ_O_xSTz870tnM08PJXQOwg9OsNNwI_HVX4S7AuW57_FzGGyRaWSuGE5SWzRS4Tur9YP0QxV4VV0Yoaoi3IAA
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: samsum
+      type: samsum
+      config: samsum
+      split: test
+    metrics:
+    - type: rouge
+      value: 33.4484
+      name: ROUGE-1
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiNTk4Yjg1YTc4YmY0MzBiZDU4ZjFhNzI4MjZkMWU1MzBlOWNlMjQ5ODMzY2YzYzRhYjJkMGUzNmI3ZjdkMzIzZSIsInZlcnNpb24iOjF9.AqS8A1OUiM0IZFBEGirv5F3Novk8lSUYSfPc3bYWLA6t-W7wgup3qA207eGbE5j9CkDWZ7QrSG1U6Z9A0sOqAA
+    - type: rouge
+      value: 10.4249
+      name: ROUGE-2
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiN2U4NjUyNTFmOGM5OTlhZDMyMTlmM2E4OWI2NGFiMDAyMGJjMzRjNWNlMGEyYWFmNTE5ZWMxM2I0ZGZmNWNmOCIsInZlcnNpb24iOjF9.SgJcHJ4qoRWXFvFiwv1PUutWktvsxQNynVPEv-GtBgxd6WI7o561ONyco5U-5tcyE_1SbSCJzz-L-R-q3cvoDA
+    - type: rouge
+      value: 24.5802
+      name: ROUGE-L
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiZmQ5MDI5MzdiNGE5NDM0MmU5OThmZTBkNjkxMzg5N2IxNGVlODdhZTZhNjg3NzFjYWEyMzA3MTQxNjMyMjRkOCIsInZlcnNpb24iOjF9.Bg5dHqCcJjmxa-xGWNR5lD9g3quX7lKkH0pjiTd2xE5WiPoLLN2c0mYa2GovdW7__WnYwhhHC7es03jmvyZbCw
+    - type: rouge
+      value: 29.8226
+      name: ROUGE-LSUM
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiNGFhOTEwNGM1MmZkNDk2ZjQ1Y2MyNjM3MGI5MGY3MWVkM2I0MjU2NWFiYmEwMjE4MTJlZWIwOGQ2MjQ3YjgzYSIsInZlcnNpb24iOjF9.W_aQKs10oXQdKEczJBGM3iiwJgb-VaXTpyA3sGof5WbhHf9vITAQA-xvynh5LgKtXQ1zjx737hnHgjEsu_Y0Cw
+    - type: loss
+      value: 4.176078796386719
+      name: loss
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiN2JhODQ5YTZkNDZkZGYyNGU2MzkxMWU5MTEwMGM2YmVjZTA5YzI5NTMxMDNhYjhlOTAxMzFiMDYwYmM0MjEzZCIsInZlcnNpb24iOjF9.OvZrPBOR5jhkoTGBgsInkH7j3_xpacXHDoT7UIXEnyXzadfBO-O-K6fjalLNZw8wSkbjHIFcL_6S_qTTxPsNAQ
+    - type: gen_len
+      value: 65.4005
+      name: gen_len
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiM2NhYjc3ZjQzNDEwYmMzOTM0ODkyZTJhZWNhNzZhYmEyZTYxMzA2YTYzMWFjOTA5ZjlhYWMzODg3NzY1ZTUwYSIsInZlcnNpb24iOjF9.vk9bgmtQFeRwdY3VXjtrJr_5wUCIeoAkI3kO0cHxhxmJo6RvUnyXiut72FuB-mlLZvqgiNkaZ-u_bh0Z3DjuCw
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: billsum
+      type: billsum
+      config: default
+      split: test
+    metrics:
+    - type: rouge
+      value: 40.5843
+      name: ROUGE-1
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiNTVjMDkyMWZjYTQ0NzgzNGUxZjNiMTg3NjU1MWJlNTQ2MWQ1NjE1MDk1OTU4ZjJiNGQ5ODg3Y2VlMWUyMzllNyIsInZlcnNpb24iOjF9.OhqBcVIuHk7fzmdrsWMvUe1bLeVMZVstZUoZpP7C1vR-3aIDl7r6eBmPrt5w-KcNq5p4teNPBsq7oKzbd5ZgDQ
+    - type: rouge
+      value: 17.3401
+      name: ROUGE-2
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiNGQxYmQzMmE0OTcyNTM5NmMwNjIxNzYxZDcwMDFkYzJkOWY4YWY3NTdhZGRhZDdlMDAxNzcwODQ5OGM3Mzc1MCIsInZlcnNpb24iOjF9.Pksn25EEqvmx757N7Swrd4yXc_xU7-AMN9yNe8lrbBa-l1LoI_2PUASvnjML4f705cfuyMAfb0FkFp5WfER2AA
+    - type: rouge
+      value: 25.1256
+      name: ROUGE-L
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiMjhjYzI5MDBiMjk2NTY3MDNmZTdiOGYwMTRlYjIwZjAwMjdlNTAyYzdhYTJlODQ4MjYzYmQ3MjRlYTA2YzhhZSIsInZlcnNpb24iOjF9.1jPepsweS2bzIqDverQzzhmhFGch7gpoEGFGqQ8zW7K10aUKWFX8lt-uZAmTa1Z5ZhzyXGBzc3dReFPhWRRJBg
+    - type: rouge
+      value: 34.6619
+      name: ROUGE-LSUM
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiM2VkZDIxNWJjOTA0NzFjOTIwOTdjYjc1M2EyNDVjZjY2ZjY3MjIxNDk3YTc5YWExNzAwN2FhOTc1NjVhYjBkYiIsInZlcnNpb24iOjF9.8opqHSUckPohoSF9jfPTpXDz2AtDwvdMqOdIXx2kE1tkOcbLPbOBfcc8RhRR98y8S26yC6EYFhFnf03CV2ejAQ
+    - type: loss
+      value: 4.792657375335693
+      name: loss
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiYTY5ZTRkMGU3OGVkODMzMDU5OWE1NTM5YjA4NDliZDlmNzc2NzZjNjFmNTA3M2EwY2NmN2E0MWJmZjQ5ZDliMiIsInZlcnNpb24iOjF9.KCKdk8xt2NWcMmYKV3-9eVEsFm9MqGllSMu9QCFJFIQlnyNXllHKdBLouoaGQz8IRYXvZKH8_TLDPIQx-31jAg
+    - type: gen_len
+      value: 163.9394
+      name: gen_len
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiYzdkZDYyZGUzYmFkZmI2NjUwYmQ0MzZjMmIyZjI1YTFiMzM4OThiZjBiMzljOTVkZTgwMjA0NTE5OGM2YmFjMiIsInZlcnNpb24iOjF9.XyMZLUdkUIF32KTJMuv_bJswQCx_Tfg4Fx823cURUixSeoIKps8_a634AreZ3Z8kb7bfE_sFGh3rM9KWsMxlDw
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: multi_news
+      type: multi_news
+      config: default
+      split: test
+    metrics:
+    - type: rouge
+      value: 39.0834
+      name: ROUGE-1
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiNjYzMmVlMDM4MTNkMTI4MjAyMTU2YTg1ZWQwNTI1MmJlNGUwZmE1NTRmYTljZTQwY2RlMjcxOTgyZGMyYTc0ZiIsInZlcnNpb24iOjF9.6yuSr7UmsFatwqQ-mEO4gmsEtWI05kGB5Ib2pnl05H1OiPT2uUwmqdUytUw8KTx9u1jv9q0cTF1cL-n2kPEJAA
+    - type: rouge
+      value: 11.4043
+      name: ROUGE-2
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiMWI5N2U2ZWI1ODM2MWUwOTIzYTAzNmRhNDA2OWEzZWRjMGEzMjBmY2EwN2YyYzU1NWE0YjIyZDE3MWE0MmMxZCIsInZlcnNpb24iOjF9.wonuxbBl25TzEaHUH_E816nHJ1OSXKfkaq7eJzbLpsfeGwcDklxUSxZxRO7VBiBMaY3Qttf9ywmEIPp40HnpBA
+    - type: rouge
+      value: 19.1813
+      name: ROUGE-L
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiZjU1NDZhN2NkMzZiZGJkODE4NDZiYjViOTZkNGMyNDlkNjBlZmFjYzU1N2IzMjFjYjY1MDU1Zjk2MzA0M2U4NyIsInZlcnNpb24iOjF9.bTCRzv3J9NiCh4aV23tAWGTvrdQCv_RS40zGwC4AJXtGS40cY7tJHYwBf9U9_rCetDBxqfjJpdaUbCAOglxLAA
+    - type: rouge
+      value: 35.1581
+      name: ROUGE-LSUM
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiMDNhNTUyZjE4NjYxYjIzYThmMDM2YWNhM2QwYzY1ODI2ZTE3NmNjMmVhOTAzZjZlOWQwYzc1NzU2NDNjNzIxMyIsInZlcnNpb24iOjF9.cWlSbEBgrMN5D-fV_yL9geNMyMkIItcVO3wehNJPzFi3E0v1-4q8pnX-UgjLzto8X7JLi6as2V_HtZE4-C-CDw
+    - type: loss
+      value: 4.654905319213867
+      name: loss
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiYTc5Nzk0ODhiNWUzNTAxNzk2YzZmMjU2NDliY2UzOTYyYTdmZGEyYjI5NDNhOTE0MGUxOTgxMGVjMmNhM2UyMSIsInZlcnNpb24iOjF9.eBBAebcl3AwkrjR6a8BvoSjDfpw8LWTRFjyIFHVzspvoOKVfnO8_NB_UeR_K127OwXyoZ70Z7X_aKJOe-2kTDA
+    - type: gen_len
+      value: 186.2494
+      name: gen_len
+      verified: true
+      verifyToken: eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJoYXNoIjoiOWI2NjVlYjgwYWJiMjcyMDUzMzEwNDNjZTMxMDM0MjAzMzk1ZmIwY2Q1ZDQ2Y2M5NDBlMDEzYzFkNWEyNzJmNiIsInZlcnNpb24iOjF9.iZ1Iy7FuWL4GH7LS5EylVj5eZRC3L2ZsbYQapAkMNzR_VXPoMGvoM69Hp-kU7gW55tmz2V4Qxhvoz9cM8fciBA
+---
+
+# Longformer Encoder-Decoder (LED) for Narrative-Esque Long Text Summarization
+
+<a href="https://colab.research.google.com/gist/pszemraj/3eba944ddc9fc9a4a1bfb21e83b57620/summarization-token-batching.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+</a>
+
+A fine-tuned version of [allenai/led-large-16384](https://huggingface.co/allenai/led-large-16384) on the `BookSum` dataset.
+
+Goal: a model that can generalize well and is useful in summarizing long text in academic and daily usage. The result works well on lots of text and can handle 16384 tokens/batch (_if you have the GPU memory to handle that_)
+
+ - See the Colab demo linked above or try the [demo on Spaces](https://huggingface.co/spaces/pszemraj/summarize-long-text)
+
+
+> Note: the API is set to generate a max of 64 tokens for runtime reasons, so the summaries may be truncated (depending on the length of input text). For best results use python as below.
+
+---
+
+# Usage - Basic
+
+- use `encoder_no_repeat_ngram_size=3` when calling the pipeline object to improve summary quality.
+  - this forces the model to use new vocabulary and create an abstractive summary, otherwise it may compile the best _extractive_ summary from the input provided.
+
+Load the model into a pipeline object:
+
+```python
+import torch
+from transformers import pipeline
+
+hf_name = 'pszemraj/led-large-book-summary'
+
+summarizer = pipeline(
+    "summarization",
+    hf_name,
+    device=0 if torch.cuda.is_available() else -1,
+)
+```
+
+- put words into the pipeline object:
+
+```python
+wall_of_text = "your words here"
+
+result = summarizer(
+    wall_of_text,
+    min_length=16,
+    max_length=256,
+    no_repeat_ngram_size=3,
+    encoder_no_repeat_ngram_size=3,
+    repetition_penalty=3.5,
+    num_beams=4,
+    early_stopping=True,
+)
+```
+
+
+**Important:** To generate the best quality summaries, you should use the global attention mask when decoding, as demonstrated in [this community notebook here](https://colab.research.google.com/drive/12INTTR6n64TzS4RrXZxMSXfrOd9Xzamo?usp=sharing), see the definition of `generate_answer(batch)`.
+
+If having computing constraints, try the base version [`pszemraj/led-base-book-summary`](https://huggingface.co/pszemraj/led-base-book-summary)
+- all the parameters for generation on the API here are the same as [the base model](https://huggingface.co/pszemraj/led-base-book-summary) for easy comparison between versions.
+
+## Training and evaluation data
+
+- the [booksum](https://arxiv.org/abs/2105.08209) dataset (this is what adds the `bsd-3-clause` license)
+- During training, the input text was the text of the `chapter`, and the output was `summary_text`
+- Eval results can be found [here](https://huggingface.co/datasets/autoevaluate/autoeval-staging-eval-project-kmfoda__booksum-79c1c0d8-10905463) with metrics on the sidebar.
+
+## Training procedure
+
+- Training completed on the BookSum dataset for 13 total epochs
+- **The final four epochs combined the training and validation sets as 'train' in an effort to increase generalization.**
+
+### Training hyperparameters
+
+#### Initial Three Epochs
+
+The following hyperparameters were used during training:
+- learning_rate: 5e-05
+- train_batch_size: 1
+- eval_batch_size: 1
+- seed: 42
+- distributed_type: multi-GPU
+- gradient_accumulation_steps: 4
+- total_train_batch_size: 4
+- optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
+- lr_scheduler_type: linear
+- num_epochs: 3
+
+#### In-between Epochs
+
+Unfortunately, don't have all records on-hand for middle epochs; the following should be representative:
+
+- learning_rate: 4e-05
+- train_batch_size: 2
+- eval_batch_size: 2
+- seed: 42
+- distributed_type: multi-GPU
+- gradient_accumulation_steps: 16
+- total_train_batch_size: 32
+- optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
+- lr_scheduler_type: cosine
+- lr_scheduler_warmup_ratio: 0.05
+- num_epochs: 6 (in addition to prior model)
+
+#### Final Two Epochs
+
+The following hyperparameters were used during training:
+- learning_rate: 2e-05
+- train_batch_size: 1
+- eval_batch_size: 1
+- seed: 42
+- distributed_type: multi-GPU
+- gradient_accumulation_steps: 16
+- total_train_batch_size: 16
+- optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
+- lr_scheduler_type: cosine
+- lr_scheduler_warmup_ratio: 0.03
+- num_epochs: 2 (in addition to prior model)
+
+
+### Framework versions
+
+- Transformers 4.19.2
+- Pytorch 1.11.0+cu113
+- Datasets 2.2.2
+- Tokenizers 0.12.1

config.json

@@ -0,0 +1,69 @@
+{
+  "_name_or_path": "pszemraj/led-large-book-summary",
+  "_num_labels": 3,
+  "activation_dropout": 0.0,
+  "activation_function": "gelu",
+  "architectures": [
+    "LEDForConditionalGeneration"
+  ],
+  "attention_dropout": 0.0,
+  "attention_window": [
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024,
+    1024
+  ],
+  "bos_token_id": 0,
+  "classif_dropout": 0.0,
+  "classifier_dropout": 0.0,
+  "d_model": 1024,
+  "decoder_attention_heads": 16,
+  "decoder_ffn_dim": 4096,
+  "decoder_layerdrop": 0.0,
+  "decoder_layers": 12,
+  "decoder_start_token_id": 2,
+  "dropout": 0.1,
+  "early_stopping": true,
+  "encoder_attention_heads": 16,
+  "encoder_ffn_dim": 4096,
+  "encoder_layerdrop": 0.0,
+  "encoder_layers": 12,
+  "eos_token_id": 2,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1",
+    "2": "LABEL_2"
+  },
+  "init_std": 0.02,
+  "is_encoder_decoder": true,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1,
+    "LABEL_2": 2
+  },
+  "length_penalty": 0.8,
+  "max_decoder_position_embeddings": 1024,
+  "max_encoder_position_embeddings": 16384,
+  "max_length": 1024,
+  "min_length": 8,
+  "model_type": "led",
+  "no_repeat_ngram_size": 3,
+  "num_beams": 4,
+  "num_hidden_layers": 12,
+  "output_past": false,
+  "pad_token_id": 1,
+  "prefix": " ",
+  "repetition_penalty": 3.5,
+  "torch_dtype": "float32",
+  "transformers_version": "4.19.2",
+  "use_cache": false,
+  "vocab_size": 50265
+}

ds_config_zero2.json

@@ -0,0 +1,39 @@
+{
+    "amp": {
+    "enabled": "auto",
+    "opt_level": "auto"
+    },
+
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "betas": "auto",
+            "eps": "auto",
+            "weight_decay": "auto"
+        }
+    },
+
+
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu",
+            "pin_memory": true
+        },
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 2e8,
+        "round_robin_gradients": true,
+        "contiguous_gradients": true
+    },
+
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 4000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

latest

@@ -0,0 +1 @@
+global_step296

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6591dc2bafd2bdca9f277d71524cb86a3a637793dc8661e08957b5dd2b52ff15
+size 135

rng_state_0.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:58f013dc79528514b3c1cce2c180b789f59730e3d4dce985927ddda73bae54d2
+size 130

special_tokens_map.json

@@ -0,0 +1 @@
+{"bos_token": {"content": "<s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "eos_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "unk_token": {"content": "<unk>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "sep_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "pad_token": {"content": "<pad>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "cls_token": {"content": "<s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "mask_token": {"content": "<mask>", "single_word": false, "lstrip": true, "rstrip": false, "normalized": true}}

tokenizer_config.json

@@ -0,0 +1 @@
+{"errors": "replace", "bos_token": {"content": "<s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "eos_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "sep_token": {"content": "</s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "cls_token": {"content": "<s>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "unk_token": {"content": "<unk>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "pad_token": {"content": "<pad>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "mask_token": {"content": "<mask>", "single_word": false, "lstrip": true, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "add_prefix_space": false, "trim_offsets": true, "model_max_length": 16384, "special_tokens_map_file": "/root/.cache/huggingface/transformers/2ad921573d53ebf0c0450d63a211e61d8e328324e84830c365abff01f2d115f1.cb2244924ab24d706b02fd7fcedaea4531566537687a539ebb94db511fd122a0", "name_or_path": "pszemraj/led-large-book-summary", "tokenizer_class": "LEDTokenizer"}

trainer_state.json

@@ -0,0 +1,364 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 2.9914638001896936,
+  "global_step": 294,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.05,
+      "learning_rate": 7.5e-06,
+      "loss": 0.481,
+      "step": 5
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 1.5e-05,
+      "loss": 0.5091,
+      "step": 10
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 2.25e-05,
+      "loss": 0.4303,
+      "step": 15
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 3e-05,
+      "loss": 0.4055,
+      "step": 20
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 2.998662940889891e-05,
+      "loss": 0.4338,
+      "step": 25
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 2.9946541471956496e-05,
+      "loss": 0.4793,
+      "step": 30
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 2.9879807655761145e-05,
+      "loss": 0.4291,
+      "step": 35
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 2.9786546929722055e-05,
+      "loss": 0.4908,
+      "step": 40
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 2.966692555397705e-05,
+      "loss": 0.463,
+      "step": 45
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 2.9521156782993066e-05,
+      "loss": 0.528,
+      "step": 50
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 2.9349500485387718e-05,
+      "loss": 0.5178,
+      "step": 55
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 2.9152262680649704e-05,
+      "loss": 0.4602,
+      "step": 60
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 2.8929794993583937e-05,
+      "loss": 0.5044,
+      "step": 65
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 2.8682494027454e-05,
+      "loss": 0.4217,
+      "step": 70
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 2.8410800656939512e-05,
+      "loss": 0.502,
+      "step": 75
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 2.811519924216873e-05,
+      "loss": 0.4549,
+      "step": 80
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 2.779621676522777e-05,
+      "loss": 0.5692,
+      "step": 85
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 2.7454421890685647e-05,
+      "loss": 0.4312,
+      "step": 90
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 2.709042395181008e-05,
+      "loss": 0.4938,
+      "step": 95
+    },
+    {
+      "epoch": 1.02,
+      "learning_rate": 2.6704871864281377e-05,
+      "loss": 0.5433,
+      "step": 100
+    },
+    {
+      "epoch": 1.07,
+      "learning_rate": 2.6298452969340952e-05,
+      "loss": 0.3459,
+      "step": 105
+    },
+    {
+      "epoch": 1.12,
+      "learning_rate": 2.58718918084368e-05,
+      "loss": 0.3739,
+      "step": 110
+    },
+    {
+      "epoch": 1.17,
+      "learning_rate": 2.5425948831550528e-05,
+      "loss": 0.375,
+      "step": 115
+    },
+    {
+      "epoch": 1.22,
+      "learning_rate": 2.496141904150859e-05,
+      "loss": 0.3809,
+      "step": 120
+    },
+    {
+      "epoch": 1.27,
+      "learning_rate": 2.447913057669456e-05,
+      "loss": 0.4183,
+      "step": 125
+    },
+    {
+      "epoch": 1.32,
+      "learning_rate": 2.3979943234689226e-05,
+      "loss": 0.4207,
+      "step": 130
+    },
+    {
+      "epoch": 1.37,
+      "learning_rate": 2.3464746939470288e-05,
+      "loss": 0.3767,
+      "step": 135
+    },
+    {
+      "epoch": 1.42,
+      "learning_rate": 2.2934460154904436e-05,
+      "loss": 0.4248,
+      "step": 140
+    },
+    {
+      "epoch": 1.48,
+      "learning_rate": 2.2390028247360042e-05,
+      "loss": 0.3374,
+      "step": 145
+    },
+    {
+      "epoch": 1.53,
+      "learning_rate": 2.183242180035951e-05,
+      "loss": 0.4582,
+      "step": 150
+    },
+    {
+      "epoch": 1.58,
+      "learning_rate": 2.1262634884275948e-05,
+      "loss": 0.4153,
+      "step": 155
+    },
+    {
+      "epoch": 1.63,
+      "learning_rate": 2.068168328415864e-05,
+      "loss": 0.409,
+      "step": 160
+    },
+    {
+      "epoch": 1.68,
+      "learning_rate": 2.0090602688846884e-05,
+      "loss": 0.4023,
+      "step": 165
+    },
+    {
+      "epoch": 1.73,
+      "learning_rate": 1.9490446844600375e-05,
+      "loss": 0.3426,
+      "step": 170
+    },
+    {
+      "epoch": 1.78,
+      "learning_rate": 1.888228567653781e-05,
+      "loss": 0.4059,
+      "step": 175
+    },
+    {
+      "epoch": 1.83,
+      "learning_rate": 1.8267203381232774e-05,
+      "loss": 0.4449,
+      "step": 180
+    },
+    {
+      "epoch": 1.88,
+      "learning_rate": 1.764629649386713e-05,
+      "loss": 0.4362,
+      "step": 185
+    },
+    {
+      "epoch": 1.93,
+      "learning_rate": 1.7020671933387917e-05,
+      "loss": 0.4874,
+      "step": 190
+    },
+    {
+      "epoch": 1.98,
+      "learning_rate": 1.63914450291526e-05,
+      "loss": 0.3326,
+      "step": 195
+    },
+    {
+      "epoch": 2.04,
+      "learning_rate": 1.5633197410233404e-05,
+      "loss": 0.4035,
+      "step": 200
+    },
+    {
+      "epoch": 2.09,
+      "learning_rate": 1.5e-05,
+      "loss": 0.3291,
+      "step": 205
+    },
+    {
+      "epoch": 2.14,
+      "learning_rate": 1.4366802589766598e-05,
+      "loss": 0.353,
+      "step": 210
+    },
+    {
+      "epoch": 2.19,
+      "learning_rate": 1.373473400935433e-05,
+      "loss": 0.319,
+      "step": 215
+    },
+    {
+      "epoch": 2.24,
+      "learning_rate": 1.3104921076168065e-05,
+      "loss": 0.341,
+      "step": 220
+    },
+    {
+      "epoch": 2.29,
+      "learning_rate": 1.247848658636778e-05,
+      "loss": 0.3276,
+      "step": 225
+    },
+    {
+      "epoch": 2.34,
+      "learning_rate": 1.185654731320877e-05,
+      "loss": 0.3628,
+      "step": 230
+    },
+    {
+      "epoch": 2.39,
+      "learning_rate": 1.124021201611919e-05,
+      "loss": 0.2727,
+      "step": 235
+    },
+    {
+      "epoch": 2.45,
+      "learning_rate": 1.0630579464064182e-05,
+      "loss": 0.3466,
+      "step": 240
+    },
+    {
+      "epoch": 2.5,
+      "learning_rate": 1.0028736476720464e-05,
+      "loss": 0.3187,
+      "step": 245
+    },
+    {
+      "epoch": 2.55,
+      "learning_rate": 9.435755986953485e-06,
+      "loss": 0.3837,
+      "step": 250
+    },
+    {
+      "epoch": 2.6,
+      "learning_rate": 8.852695128051192e-06,
+      "loss": 0.2955,
+      "step": 255
+    },
+    {
+      "epoch": 2.65,
+      "learning_rate": 8.280593349124432e-06,
+      "loss": 0.3793,
+      "step": 260
+    },
+    {
+      "epoch": 2.7,
+      "learning_rate": 7.720470562033787e-06,
+      "loss": 0.3443,
+      "step": 265
+    },
+    {
+      "epoch": 2.75,
+      "learning_rate": 7.17332532314626e-06,
+      "loss": 0.2915,
+      "step": 270
+    },
+    {
+      "epoch": 2.8,
+      "learning_rate": 6.640133053163455e-06,
+      "loss": 0.3514,
+      "step": 275
+    },
+    {
+      "epoch": 2.85,
+      "learning_rate": 6.12184429819474e-06,
+      "loss": 0.3221,
+      "step": 280
+    },
+    {
+      "epoch": 2.9,
+      "learning_rate": 5.619383035175448e-06,
+      "loss": 0.2903,
+      "step": 285
+    },
+    {
+      "epoch": 2.95,
+      "learning_rate": 5.133645024651171e-06,
+      "loss": 0.3397,
+      "step": 290
+    }
+  ],
+  "max_steps": 392,
+  "num_train_epochs": 4,
+  "total_flos": 1.821325738775675e+17,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:401745ea2f5f8bc223616c82c7bae7b501398125eba6ecd872c79ffcc7abc266
+size 129

zero_to_fp32.py

@@ -0,0 +1,484 @@
+#!/usr/bin/env python
+
+# This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example: python zero_to_fp32.py . pytorch_model.bin
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+from collections import OrderedDict
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+import deepspeed
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION,
+                                            OPTIMIZER_STATE_DICT,
+                                            PARAM_SHAPES,
+                                            SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS,
+                                            ZERO_STAGE,
+                                            PARTITION_COUNT,
+                                            PARAM_SHAPES,
+                                            BUFFER_NAMES)
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage == 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_optim_files(checkpoint_dir):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    optim_files = sorted(glob.glob(os.path.join(checkpoint_dir,
+                                                "*_optim_states.pt")),
+                         key=natural_keys)
+
+    if len(optim_files) == 0:
+        raise FileNotFoundError(
+            f"can't find '*_optim_states.pt' files in directory '{checkpoint_dir}'")
+
+    return optim_files
+
+
+def parse_model_state(file):
+    state_dict = torch.load(file, map_location=device)
+
+    if BUFFER_NAMES not in state_dict:
+        raise ValueError(f"{file} is not a model state checkpoint")
+    buffer_names = state_dict[BUFFER_NAMES]
+    if debug:
+        print("Found buffers:", buffer_names)
+
+    # recover just the buffers while restoring them to fp32 if they were saved in fp16
+    buffers = {
+        k: v.float()
+        for k,
+        v in state_dict["module"].items() if k in buffer_names
+    }
+    param_shapes = state_dict[PARAM_SHAPES]
+
+    ds_version = state_dict.get(DS_VERSION, None)
+
+    return buffers, param_shapes, ds_version
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+
+    total_files = len(files)
+    state_dicts = []
+    for f in files:
+        state_dicts.append(torch.load(f, map_location=device))
+
+    if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage == 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    if zero_stage == 2:
+        fp32_flat_groups = [
+            state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key]
+            for i in range(len(state_dicts))
+        ]
+    elif zero_stage == 3:
+        # if there is more than one param group, there will be multiple flattened tensors - one
+        # flattened tensor per group - for simplicity merge them into a single tensor
+        #
+        # XXX: could make the script more memory efficient for when there are multiple groups - it
+        # will require matching the sub-lists of param_shapes for each param group flattened tensor
+
+        fp32_flat_groups = [
+            torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key],
+                      0) for i in range(len(state_dicts))
+        ]
+
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(
+        f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_file = get_model_state_file(ds_checkpoint_dir, zero_stage)
+    buffers, param_shapes, ds_version = parse_model_state(model_file)
+    print(f'Parsing checkpoint created by deepspeed=={ds_version}')
+
+    if zero_stage == 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size,
+                                                          param_shapes,
+                                                          fp32_flat_groups,
+                                                          buffers)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size,
+                                                          param_shapes,
+                                                          fp32_flat_groups,
+                                                          buffers)
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size,
+                                               param_shapes,
+                                               fp32_flat_groups,
+                                               buffers):
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(
+                    f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum([
+        full_single_fp32_vector.numel()
+        for full_single_fp32_vector in merged_single_partition_of_fp32_groups
+    ])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum(
+            [sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    state_dict = OrderedDict()
+
+    # buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel()
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(
+                    f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} "
+                )
+            state_dict[name] = full_single_fp32_vector.narrow(
+                0,
+                offset,
+                unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(
+                f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(
+        f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements"
+    )
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size,
+                                               param_shapes,
+                                               fp32_flat_groups,
+                                               buffers):
+
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    avail_numel = fp32_flat_groups[0].numel() * world_size
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    state_dict = OrderedDict()
+
+    # buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    for name, shape in param_shapes.items():
+
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"{total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # XXX: memory usage doubles here
+        state_dict[name] = torch.cat(
+            tuple(fp32_flat_groups[i].narrow(0,
+                                             offset,
+                                             partitioned_numel)
+                  for i in range(world_size)),
+            0).narrow(0,
+                      0,
+                      unpartitioned_numel).view(shape)
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(
+            f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(
+        f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements"
+    )
+
+    return state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
+    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+    """
+
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+    print(f"Saving fp32 state dict to {output_file}")
+    torch.save(state_dict, output_file)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info(f"Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info(f"Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "checkpoint_dir",
+        type=str,
+        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help=
+        "path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)"
+    )
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)