Update README.md

README.md

@@ -1,13 +1,16 @@
 ---
-language: en
+language: 
+- en
+- zh
 tags:
-- SEGA
+- GENIUS
+- conditional text generation
+- sketch-based text generation
 - data augmentation
-- keywords-to-text generation
-- sketch-to-text generation
 license: apache-2.0
 datasets:
 - c4
+- beyond/chinese_clean_passages_80m
 
 
 widget:
@@ -28,48 +31,59 @@ inference:
     num_beams: 3
     do_sample: True
 ---
+# 💡GENIUS – generating text using sketches!
 
-# SEGA-large model
+- **Paper: [GENIUS: Sketch-based Language Model Pre-training via Extreme and Selective Masking for Text Generation and Augmentation](https://github.com/beyondguo/genius/blob/master/GENIUS_gby_arxiv.pdf)**
+- **GitHub: [GENIUS project, GENIUS pre-training, GeniusAug for data augmentation](https://github.com/beyondguo/genius)**
 
-**SEGA: SkEtch-based Generative Augmentation** \
-**基于草稿的生成式增强模型**
+💡**GENIUS** is a powerful conditional text generation model using sketches as input, which can fill in the missing contexts for a given **sketch** (key information consisting of textual spans, phrases, or words, concatenated by mask tokens). GENIUS is pre-trained on a large-scale textual corpus with a novel *reconstruction from sketch* objective using an *extreme and selective masking* strategy, enabling it to generate diverse and high-quality texts given sketches.
 
-**SEGA** is a **general text augmentation model** that can be used for data augmentation for **various NLP tasks** (including sentiment analysis, topic classification, NER, and QA). SEGA uses an encoder-decoder structure (based on the BART architecture) and is pre-trained on the `C4-realnewslike` corpus. 
 
+**GENIUS** can also be used as a general textual **data augmentation tool** for **various NLP tasks** (including sentiment analysis, topic classification, NER, and QA). 
 
-![sega-illustration](https://cdn.jsdelivr.net/gh/beyondguo/mdnice_pictures/typora/sega-main-illustration.png)
 
-- Paper: [SEGA: SkEtch-based Generative Augmentation (preprint)](https://github.com/beyondguo/SEGA/blob/master/SEGA_gby_preprint.pdf)
-- GitHub: [SEGA](https://github.com/beyondguo/SEGA). 
+![image-20221119164544165](https://cdn.jsdelivr.net/gh/beyondguo/mdnice_pictures/typora/hi-genius.png)
 
-**SEGA** is able to write complete paragraphs given a *sketch*, which can be composed of:
-- keywords /key-phrases, like "––NLP––AI––computer––science––"
-- spans, like "Conference on Empirical Methods––submission of research papers––"
-- sentences, like "I really like machine learning––I work at Google since last year––"
-- or mixup~
 
+- Models hosted in 🤗 Huggingface:
 
 **Model variations:**
+
 | Model | #params | Language | comment|
 |------------------------|--------------------------------|-------|---------|
-| [`sega-large`](https://huggingface.co/beyond/sega-large) | 406M   | English | The version used in paper |
-| [`sega-large-k2t`](https://huggingface.co/beyond/sega-large-k2t)  | 406M    | English | keywords-to-text |
-| [`sega-base`](https://huggingface.co/beyond/sega-base)  | 139M    | English | smaller version |
-| [`sega-base-ps`](https://huggingface.co/beyond/sega-base)  | 139M    | English | pre-trained both in paragraphs and short sentences |
-| [`sega-base-chinese`](https://huggingface.co/beyond/sega-base-chinese) | 116M    | 中文 | 在一千万纯净中文段落上预训练|
+| [`genius-large`](https://huggingface.co/beyond/genius-large) | 406M   | English | The version used in **paper** (recommend) |
+| [`genius-large-k2t`](https://huggingface.co/beyond/genius-large-k2t)  | 406M    | English | keywords-to-text |
+| [`genius-base`](https://huggingface.co/beyond/genius-base)  | 139M    | English | smaller version |
+| [`genius-base-ps`](https://huggingface.co/beyond/genius-base)  | 139M    | English | pre-trained both in paragraphs and short sentences |
+| [`genius-base-chinese`](https://huggingface.co/beyond/genius-base-chinese) | 116M    | 中文 | 在一千万纯净中文段落上预训练|
 
 ---
 
-### How to use
-#### 1. If you want to generate sentences given a **sketch**
+## Usage
+
+### What is a sketch?
+
+First, what is a **sketch**? As defined in our paper, a sketch is "key information consisting of textual spans, phrases, or words, concatenated by mask tokens". It's like a draft or framework when you begin to write an article. With GENIUS model, you can input some key elements you want to mention in your wrinting, then the GENIUS model can generate cohrent text based on your sketch.
+
+The sketch which can be composed of:
+
+- keywords /key-phrases, like `__NLP__AI__computer__science__`
+- spans, like `Conference on Empirical Methods__submission of research papers__`
+- sentences, like `I really like machine learning__I work at Google since last year__`
+- or a mixup!
+
+(the `__` is the mask token. Use `<mask>` for English, and `[MASK]` for Chinese)
+
+### How to use the model
+#### 1. If you already have a sketch in mind, and want to get a paragraph based on it...
 ```python
 from transformers import pipeline
 # 1. load the model with the huggingface `pipeline`
-sega = pipeline("text2text-generation", model='beyond/sega-large', device=0)
+genius = pipeline("text2text-generation", model='beyond/genius-large', device=0)
 # 2. provide a sketch (joint by <mask> tokens)
 sketch = "<mask> Conference on Empirical Methods <mask> submission of research papers <mask> Deep Learning <mask>"
-# 3. just do it!
-generated_text = sega(sketch, num_beams=3, do_sample=True, max_length=200)[0]['generated_text']
+# 3. here we go!
+generated_text = genius(sketch, num_beams=3, do_sample=True, max_length=200)[0]['generated_text']
 print(generated_text)
 ```
 Output:
@@ -77,14 +91,20 @@ Output:
 'The Conference on Empirical Methods welcomes the submission of research papers. Abstracts should be in the form of a paper or presentation. Please submit abstracts to the following email address: eemml.stanford.edu. The conference will be held at Stanford University on April 1618, 2019. The theme of the conference is Deep Learning.'
 ```
 
-#### 2. If you want to do **data augmentation** to generate new training samples
-Please Check our Github page: [github.com/beyondguo/SEGA](https://github.com/beyondguo/SEGA), where we provide ready-to-run scripts for data augmentation for text classification/NER/MRC tasks.
+If you have a lot of sketches, you can batch-up your sketches to a Huggingface `Dataset` object, which can be much faster.
+
+TODO: we are also building a python package for more convenient use of GENIUS, which will be released in few weeks.
+
+#### 2. If you have an NLP dataset (e.g. classification) and want to do data augmentation to enlarge your dataset...
+
+Please check [genius/augmentation_clf](https://github.com/beyondguo/genius/tree/master/augmentation_clf) and [genius/augmentation_ner_qa](https://github.com/beyondguo/genius/tree/master/augmentation_ner_qa), where we provide ready-to-run scripts for data augmentation for text classification/NER/MRC tasks.
 
----
 
-## SEGA as A Strong Data Augmentation Tool:
+
+## Augmentation Experiments:
+Data augmentation is an important application for natural language generation (NLG) models, which is also a valuable evaluation of whether the generated text can be used in real applications. 
 - Setting: Low-resource setting, where only n={50,100,200,500,1000} labeled samples are available for training. The below results are the average of all training sizes.
-- Datasets: [HuffPost](https://huggingface.co/datasets/khalidalt/HuffPost), [BBC](https://huggingface.co/datasets/SetFit/bbc-news), [SST2](https://huggingface.co/datasets/glue), [IMDB](https://huggingface.co/datasets/imdb), [Yahoo](https://huggingface.co/datasets/yahoo_answers_topics), [20NG](https://huggingface.co/datasets/newsgroup).
+- Text Classification Datasets: [HuffPost](https://huggingface.co/datasets/khalidalt/HuffPost), [BBC](https://huggingface.co/datasets/SetFit/bbc-news), [SST2](https://huggingface.co/datasets/glue), [IMDB](https://huggingface.co/datasets/imdb), [Yahoo](https://huggingface.co/datasets/yahoo_answers_topics), [20NG](https://huggingface.co/datasets/newsgroup).
 - Base classifier: [DistilBERT](https://huggingface.co/distilbert-base-cased)
 
 
@@ -98,8 +118,8 @@ In-distribution (ID) evaluations:
 |    C-MLM   |   80.60   |   96.13   |   45.40   |   46.36   |   77.31   |   76.91   |   70.45   |
 |   LAMBADA  |   81.46   |   93.74   |   50.49   |   47.72   |   78.22   |   78.31   |   71.66   |
 |     STA    |   80.74   |   95.64   |   46.96   |   47.27   |   77.88   |   77.80   |   71.05   |
-|  **SEGA**  |   81.43   |   95.74   |   49.60   |   50.38   | **80.16** |   78.82   |   72.68   |
-| **SEGA-f** | **81.82** |   95.99   | **50.42** | **50.81** |   79.40   | **80.57** | **73.17** |
+|  **GeniusAug**  |   81.43   |   95.74   |   49.60   |   50.38   | **80.16** |   78.82   |   72.68   |
+| **GeniusAug-f** | **81.82** |   95.99   | **50.42** | **50.81** |   79.40   | **80.57** | **73.17** |
 
 Out-of-distribution (OOD) evaluations:
 |            |  Huff->BBC |  BBC->Huff | IMDB->SST2 | SST2->IMDB |    avg.    |
@@ -111,10 +131,11 @@ Out-of-distribution (OOD) evaluations:
 |    C-MLM   |   64.94   | **67.80** |   74.98   |   71.78   |   69.87   |
 |   LAMBADA  |   68.57   |   52.79   |   75.24   |   76.04   |   68.16   |
 |     STA    |   69.31   |   64.82   |   74.72   |   73.62   |   70.61   |
-|  **SEGA**  |   74.87   |   66.85   |   76.02   |   74.76   |   73.13   |
-| **SEGA-f** | **76.18** |   66.89   | **77.45** | **80.36** | **75.22** |
+|  **GeniusAug**  |   74.87   |   66.85   |   76.02   |   74.76   |   73.13   |
+| **GeniusAug-f** | **76.18** |   66.89   | **77.45** | **80.36** | **75.22** |
 
+### BibTeX entry and citation info
+TBD
 
 
-### BibTeX entry and citation info