Add app

Pipfile

@@ -0,0 +1,24 @@
+[[source]]
+url = "https://pypi.org/simple"
+verify_ssl = true
+name = "pypi"
+
+[packages]
+streamlit = "==0.84.2"
+arabic-reshaper = "==2.1.3"
+python-bidi = "==0.4.2"
+pyarabic = "*"
+farasapy = "==0.0.14"
+emoji = "==1.4.2"
+awesome-streamlit = "*"
+torch = "==1.9.0"
+transformers = "==4.10.0"
+psutil = "==5.8.0"
+fuzzysearch = "==0.7.3"
+more-itertools = "==8.9.0"
+cookiecutter = "*"
+
+[dev-packages]
+
+[requires]
+python_version = "3.8"

README.md

@@ -1,13 +1,9 @@
 ---
-title: Design Project
-emoji: 🏃
+title: Arabic NLP Demo
+emoji: ⌨
 colorFrom: purple
 colorTo: green
 sdk: streamlit
-sdk_version: 1.9.0
 app_file: app.py
-pinned: false
-license: unlicense
+pinned: true
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

app.py

@@ -0,0 +1,48 @@
+import awesome_streamlit as ast
+import streamlit as st
+
+from backend.utils import get_current_ram_usage, ga
+
+import backend.aragpt
+import backend.home
+import backend.processor
+import backend.sa
+import backend.qa
+
+st.set_page_config(
+    page_title="TEST", page_icon="📖", initial_sidebar_state="expanded", layout="wide"
+)
+
+ga(st.__file__)
+
+PAGES = {
+    "Home": backend.home,
+    "Arabic Text Preprocessor": backend.processor,
+    "Arabic Language Generation": backend.aragpt,
+    "Arabic Sentiment Analysis": backend.sa,
+    # "Arabic Sarcasm Detection": backend.sarcasm,
+    "Arabic Question Answering": backend.qa,
+}
+
+
+st.sidebar.title("Navigation")
+selection = st.sidebar.radio("Pages", list(PAGES.keys()))
+
+page = PAGES[selection]
+# with st.spinner(f"Loading {selection} ..."):
+ast.shared.components.write_page(page)
+
+st.sidebar.header("Info")
+st.sidebar.write("Made by [Wissam Antoun](https://twitter.com/wissam_antoun)")
+st.sidebar.write(
+    "Pre-trained models are available on [HF Hub](https://huggingface.co/aubmindlab)"
+)
+st.sidebar.write(
+    "Models source code available on [GitHub](https://github.com/aub-mind/arabert)"
+)
+st.sidebar.write(
+    "App source code available on [GitHub](https://github.com/WissamAntoun/Arabic-NLP-app)"
+)
+if st.sidebar.checkbox("Show RAM usage"):
+    ram = get_current_ram_usage()
+    st.sidebar.write("Ram usage: {:.2f}/{:.2f} GB".format(ram[0], ram[1]))

backend/aragpt.py

@@ -0,0 +1,189 @@
+import streamlit as st
+from .services import TextGeneration
+from tokenizers import Tokenizer
+from functools import lru_cache
+
+# @st.cache(allow_output_mutation=False, hash_funcs={Tokenizer: str})
+@lru_cache(maxsize=1)
+def load_text_generator():
+    generator = TextGeneration()
+    generator.load()
+    return generator
+
+
+generator = load_text_generator()
+
+qa_prompt = """
+            أجب عن السؤال التالي:
+            """
+qa_prompt_post = """ الجواب هو  """
+qa_prompt_post_year = """ في سنة: """
+
+
+def write():
+    st.markdown(
+        """
+        <h1 style="text-align:left;">Arabic Language Generation</h1>
+        """,
+        unsafe_allow_html=True,
+    )
+
+    # Sidebar
+
+    # Taken from https://huggingface.co/spaces/flax-community/spanish-gpt2/blob/main/app.py
+    st.sidebar.subheader("Configurable parameters")
+
+    model_name = st.sidebar.selectbox(
+        "Model Selector",
+        options=[
+            "AraGPT2-Base",
+            # "AraGPT2-Medium",
+            # "Aragpt2-Large",
+            "AraGPT2-Mega",
+        ],
+        index=0,
+    )
+
+    max_new_tokens = st.sidebar.number_input(
+        "Maximum length",
+        min_value=0,
+        max_value=1024,
+        value=100,
+        help="The maximum length of the sequence to be generated.",
+    )
+    temp = st.sidebar.slider(
+        "Temperature",
+        value=1.0,
+        min_value=0.1,
+        max_value=100.0,
+        help="The value used to module the next token probabilities.",
+    )
+    top_k = st.sidebar.number_input(
+        "Top k",
+        value=10,
+        help="The number of highest probability vocabulary tokens to keep for top-k-filtering.",
+    )
+    top_p = st.sidebar.number_input(
+        "Top p",
+        value=0.95,
+        help=" If set to float < 1, only the most probable tokens with probabilities that add up to top_p or higher are kept for generation.",
+    )
+    do_sample = st.sidebar.selectbox(
+        "Sampling?",
+        (True, False),
+        help="Whether or not to use sampling; use greedy decoding otherwise.",
+    )
+    num_beams = st.sidebar.number_input(
+        "Number of beams",
+        min_value=1,
+        max_value=10,
+        value=3,
+        help="The number of beams to use for beam search.",
+    )
+    repetition_penalty = st.sidebar.number_input(
+        "Repetition Penalty",
+        min_value=0.0,
+        value=3.0,
+        step=0.1,
+        help="The parameter for repetition penalty. 1.0 means no penalty",
+    )
+    no_repeat_ngram_size = st.sidebar.number_input(
+        "No Repeat N-Gram Size",
+        min_value=0,
+        value=3,
+        help="If set to int > 0, all ngrams of that size can only occur once.",
+    )
+
+    st.write("#")
+
+    col = st.columns(2)
+
+    col[0].image("images/AraGPT2.png", width=200)
+
+    st.markdown(
+        """
+
+        <h3 style="text-align:left;">AraGPT2 is GPT2 model trained from scratch on 77GB of Arabic text.</h3>
+        <h4 style="text-align:left;"> More details in our <a href="https://github.com/aub-mind/arabert/tree/master/aragpt2">repo</a>.</h4>
+
+        <p style="text-align:left;"><p>
+        <p style="text-align:left;">Use the generation paramters on the sidebar to adjust generation quality.</p>
+        <p style="text-align:right;"><p>
+        """,
+        unsafe_allow_html=True,
+    )
+
+    # col[0].write(
+    #     "AraGPT2 is trained from screatch on 77GB of Arabic text. More details in our [repo](https://github.com/aub-mind/arabert/tree/master/aragpt2)."
+    # )
+    # st.write("## Generate Arabic Text")
+
+    st.markdown(
+        """
+        <style>
+        p, div, input, label, textarea{
+        text-align: right;
+        }
+        </style>
+        """,
+        unsafe_allow_html=True,
+    )
+
+    prompt = st.text_area(
+        "Prompt",
+        "يحكى أن مزارعا مخادعا قام ببيع بئر الماء الموجود في أرضه لجاره مقابل مبلغ كبير من المال",
+    )
+    if st.button("Generate"):
+        with st.spinner("Generating..."):
+            generated_text = generator.generate(
+                prompt=prompt,
+                model_name=model_name,
+                max_new_tokens=max_new_tokens,
+                temperature=temp,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                do_sample=do_sample,
+                num_beams=num_beams,
+                no_repeat_ngram_size=no_repeat_ngram_size,
+            )
+            st.write(generated_text)
+
+    st.markdown("---")
+    st.subheader("")
+    st.markdown(
+        """
+        <p style="text-align:left;"><p>
+        <h2 style="text-align:left;">Zero-Shot Question Answering</h2>
+
+        <p style="text-align:left;">Adjust the maximum length to closely match the expected output length. Setting the Sampling paramter to False is recommended</p>
+        <p style="text-align:left;"><p>
+        """,
+        unsafe_allow_html=True,
+    )
+
+    question = st.text_input(
+        "Question", "من كان رئيس ألمانيا النازية في الحرب العالمية الثانية ؟"
+    )
+    is_date = st.checkbox("Help the model: Is the answer a date?")
+    if st.button("Answer"):
+
+        prompt2 = qa_prompt + question + qa_prompt_post
+        if is_date:
+            prompt2 += qa_prompt_post_year
+        else:
+            prompt2 += " : "
+        with st.spinner("Thinking..."):
+            answer = generator.generate(
+                prompt=prompt2,
+                model_name=model_name,
+                max_new_tokens=max_new_tokens,
+                temperature=temp,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                do_sample=do_sample,
+                num_beams=num_beams,
+                no_repeat_ngram_size=no_repeat_ngram_size,
+            )
+            st.write(answer)

backend/home.py

@@ -0,0 +1,156 @@
+import streamlit as st
+import awesome_streamlit as ast
+
+
+def write():
+    st.markdown(
+        """
+    # Arabic Natural Language Processing
+
+    ![visitors](https://visitor-badge.glitch.me/badge?page_id=wissamantoun.arabicnlpapp)
+
+
+    In this HuggingFace space you will be able to test the different Arabic NLP models that my colleges at [AUB MIND Lab](https://sites.aub.edu.lb/mindlab/) have built, with some other applications.
+
+    Check the **Navigation bar** to access the apps:
+    - Arabic Text Preprocessor: Test how text imput is treated by our preprocessor
+    - Arabic Language Generation: Generate Arabic text using our AraGPT2 language models
+    - Arabic Sentiment Analysis: Test the senitment analysis model that won the [Arabic Senitment Analysis competition @ KAUST](https://www.kaggle.com/c/arabic-sentiment-analysis-2021-kaust)
+    - Arabic Question Answering: Test our AraELECTRA QA capabilities
+    """
+    )
+    st.markdown("#")
+    col1, col2, col3 = st.columns(3)
+
+    col1.write("## **AraBERT**")
+    col1.image("images/arabert_logo.png", width=200)
+
+    col2.write("## **AraGPT2**")
+    col2.image("images/AraGPT2.png", width=200)
+
+    col3.write("## **AraElectra**")
+    col3.image("images/AraELECTRA.png", width=200)
+
+    st.markdown(
+        """
+
+        You can find the more details in the source code and paper linked in our repository on GitHub [repo](https://github.com/aub-mind/arabert).
+
+        ## Dataset
+
+        The pretraining data used for the new **AraBERT** model is also used for **AraGPT2 and AraELECTRA**.
+
+        The dataset consists of 77GB or 200,095,961 lines or 8,655,948,860 words or 82,232,988,358 chars (before applying Farasa Segmentation)
+
+        Our large models were train a TPUv3-128 provided by TFRC.
+
+        For the new dataset we added the unshuffled OSCAR corpus, after we thoroughly filter it, to the previous dataset used in AraBERTv1 but with out the websites that we previously crawled:
+        - OSCAR unshuffled and filtered.
+        - [Arabic Wikipedia dump](https://archive.org/details/arwiki-20190201) from 2020/09/01
+        - [The 1.5B words Arabic Corpus](https://www.semanticscholar.org/paper/1.5-billion-words-Arabic-Corpus-El-Khair/f3eeef4afb81223df96575adadf808fe7fe440b4)
+        - [The OSIAN Corpus](https://www.aclweb.org/anthology/W19-4619)
+        - Assafir news articles. Huge thank you for Assafir for the data
+
+        ## Models
+
+        Model | HuggingFace Model Name | Size (MB/Params)| Pre-Segmentation |  Hardware | Sequence Length | Batch Size | Num of Steps | Total Time (in Days) |
+        ---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:
+        AraBERTv0.2-base | [bert-base-arabertv02](https://huggingface.co/aubmindlab/bert-base-arabertv02) | 543MB / 136M | No | TPUv3-8 | 128 /512 | 2560/384 | 1M/ 2M | 36 |
+        AraBERTv0.2-large| [bert-large-arabertv02](https://huggingface.co/aubmindlab/bert-large-arabertv02) | 1.38G / 371M | No | TPUv3-128 | 128 /512 | 13440 / 2056 | 250K / 300K | 7 |
+        AraBERTv2-base| [bert-base-arabertv2](https://huggingface.co/aubmindlab/bert-base-arabertv2) | 543MB / 136M | Yes | TPUv3-8 |128 /512 | 2560 / 384 | 1M / 2M | 36 |
+        AraBERTv2-large| [bert-large-arabertv2](https://huggingface.co/aubmindlab/bert-large-arabertv2) | 1.38G / 371M | Yes | TPUv3-128 |128 /512 | 13440 / 2056|  250K / 300K | 7 |
+        AraBERTv0.1-base| [bert-base-arabertv01](https://huggingface.co/aubmindlab/bert-base-arabertv01) | 543MB / 136M | No | TPUv2-8 |128 /512 |128 / 512 | 900K / 300K| 4 |
+        AraBERTv1-base| [bert-base-arabert](https://huggingface.co/aubmindlab/bert-base-arabert) | 543MB / 136M | Yes | TPUv2-8 |128 /512 |128 / 512 | 900K / 300K| 4 |
+        AraGPT2-base | [aragpt2-base](https://huggingface.co/aubmindlab/aragpt2-base) | 527MB/135M | No | TPUv3-128 | 1024 | 1792 | 125K | 1.5 |
+        AraGPT2-medium | [aragpt2-medium](https://huggingface.co/aubmindlab/aragpt2-medium) |  1.38G/370M  | No |TPUv3-8 | 1024 | 80 | 1M | 15 |
+        AraGPT2-large | [aragpt2-large](https://huggingface.co/aubmindlab/aragpt2-large) |  2.98GB/792M  | No |TPUv3-128 | 1024 | 256 | 220k | 3 |
+        AraGPT2-mega | [aragpt2-mega](https://huggingface.co/aubmindlab/aragpt2-mega) |  5.5GB/1.46B  |No |TPUv3-128 | 1024 | 256 | 800K | 9 |
+        AraELECTRA-base-generator | [araelectra-base-generator](https://huggingface.co/aubmindlab/araelectra-base-generator) |  227MB/60M  | No | TPUv3-8 | 512 | 256 | 2M | 24
+        AraELECTRA-base-discriminator | [araelectra-base-discriminator](https://huggingface.co/aubmindlab/araelectra-base-discriminator) |  516MB/135M  | No | TPUv3-8 | 512 | 256 | 2M | 24
+        AraBERTv0.2-Twitter-base| [bert-base-arabertv02-twitter](https://huggingface.co/aubmindlab/bert-base-arabertv02-twitter) | 543MB / 136M | No | V100 | *64* | - | - | - |
+        AraBERTv0.2-Twitter-large| [bert-large-arabertv02-twitter](https://huggingface.co/aubmindlab/bert-large-arabertv02-twitter) | 1.38G / 371M | No | V100 | *64* | - | - | - |
+
+        All models are available in the `HuggingFace` model page under the [aubmindlab](https://huggingface.co/aubmindlab/) name. Checkpoints are available in PyTorch, TF2 and TF1 formats.
+
+        # Preprocessing
+
+        You can test the Arabic Preprocessing pipeline in the Arabic Text Preprocessing page.
+
+        It is recommended to apply our preprocessing function before training/testing on any dataset.
+        **Install farasapy to segment text for AraBERT v1 & v2 `pip install farasapy`**
+
+        ```python
+        from arabert.preprocess import ArabertPreprocessor
+
+        model_name = "aubmindlab/bert-base-arabertv2"
+        arabert_prep = ArabertPreprocessor(model_name=model_name)
+
+        text = "ولن نبالغ إذا قلنا: إن 'هاتف' أو 'كمبيوتر المكتب' في زمننا هذا ضروري"
+        arabert_prep.preprocess(text)
+        >>>"و+ لن نبالغ إذا قل +نا : إن ' هاتف ' أو ' كمبيوتر ال+ مكتب ' في زمن +نا هذا ضروري"
+        ```
+
+        You can also use the `unpreprocess()` function to reverse the preprocessing changes, by fixing the spacing around non alphabetical characters, and also de-segmenting if the model selected need pre-segmentation. We highly recommend unprocessing generated content of `AraGPT2` model, to make it look more natural.
+        ```python
+        output_text = "و+ لن نبالغ إذا قل +نا : إن ' هاتف ' أو ' كمبيوتر ال+ مكتب ' في زمن +نا هذا ضروري"
+        arabert_prep.unpreprocess(output_text)
+        >>>"ولن نبالغ إذا قلنا: إن 'هاتف' أو 'كمبيوتر المكتب' في زمننا هذا ضروري"
+        ```
+
+        # If you used this model please cite us as :
+
+        ## AraBERT
+        Google Scholar has our Bibtex wrong (missing name), use this instead
+        ```
+        @inproceedings{antoun2020arabert,
+        title={AraBERT: Transformer-based Model for Arabic Language Understanding},
+        author={Antoun, Wissam and Baly, Fady and Hajj, Hazem},
+        booktitle={LREC 2020 Workshop Language Resources and Evaluation Conference 11--16 May 2020},
+        pages={9}
+        }
+        ```
+        ## AraGPT2
+        ```
+        @inproceedings{antoun-etal-2021-aragpt2,
+            title = "{A}ra{GPT}2: Pre-Trained Transformer for {A}rabic Language Generation",
+            author = "Antoun, Wissam  and
+            Baly, Fady  and
+            Hajj, Hazem",
+            booktitle = "Proceedings of the Sixth Arabic Natural Language Processing Workshop",
+            month = apr,
+            year = "2021",
+            address = "Kyiv, Ukraine (Virtual)",
+            publisher = "Association for Computational Linguistics",
+            url = "https://www.aclweb.org/anthology/2021.wanlp-1.21",
+            pages = "196--207",
+        }
+        ```
+
+        ## AraELECTRA
+        ```
+        @inproceedings{antoun-etal-2021-araelectra,
+            title = "{A}ra{ELECTRA}: Pre-Training Text Discriminators for {A}rabic Language Understanding",
+            author = "Antoun, Wissam  and
+            Baly, Fady  and
+            Hajj, Hazem",
+            booktitle = "Proceedings of the Sixth Arabic Natural Language Processing Workshop",
+            month = apr,
+            year = "2021",
+            address = "Kyiv, Ukraine (Virtual)",
+            publisher = "Association for Computational Linguistics",
+            url = "https://www.aclweb.org/anthology/2021.wanlp-1.20",
+            pages = "191--195",
+        }
+        ```
+
+
+        # Acknowledgments
+        Thanks to TensorFlow Research Cloud (TFRC) for the free access to Cloud TPUs, couldn't have done it without this program, and to the [AUB MIND Lab](https://sites.aub.edu.lb/mindlab/) Members for the continous support. Also thanks to [Yakshof](https://www.yakshof.com/#/) and Assafir for data and storage access. Another thanks for Habib Rahal (https://www.behance.net/rahalhabib), for putting a face to AraBERT.
+
+        # Contacts
+        **Wissam Antoun**: [Linkedin](https://www.linkedin.com/in/wissam-antoun-622142b4/) | [Twitter](https://twitter.com/wissam_antoun) | [Github](https://github.com/WissamAntoun) | wfa07 (AT) mail (DOT) aub (DOT) edu | wissam.antoun (AT) gmail (DOT) com
+
+        **Fady Baly**: [Linkedin](https://www.linkedin.com/in/fadybaly/) | [Twitter](https://twitter.com/fadybaly) | [Github](https://github.com/fadybaly) | fgb06 (AT) mail (DOT) aub (DOT) edu | baly.fady (AT) gmail (DOT) com
+
+        """
+    )

backend/modeling_gpt2.py

@@ -0,0 +1,1599 @@
+# coding=utf-8
+# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+PyTorch OpenAI GPT-2 model.
+Adapted from https://github.com/huggingface/transformers/blob/v4.0.1/src/transformers/models/gpt2/modeling_gpt2.py
+and https://github.com/ghosthamlet/gpt2-ml-torch/blob/master/gpt2_ml_torch/modeling_gpt2.py
+"""
+
+
+import logging
+import os
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss, MSELoss
+from transformers import CONFIG_NAME, WEIGHTS_NAME, GPT2Config, GPT2Model
+from transformers.activations import ACT2FN
+from transformers.file_utils import (
+    ModelOutput,
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    replace_return_docstrings,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import (
+    Conv1D,
+    PreTrainedModel,
+    SequenceSummary,
+    find_pruneable_heads_and_indices,
+    prune_conv1d_layer,
+)
+from transformers.utils.model_parallel_utils import assert_device_map, get_device_map
+
+# THe Difference from Transformers is code under _USE_GROVER
+_USE_GROVER = True
+
+logger = logging.getLogger(__name__)
+
+_CONFIG_FOR_DOC = "GPT2Config"
+_TOKENIZER_FOR_DOC = "GPT2Tokenizer"
+
+GPT2_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "gpt2",
+    "gpt2-medium",
+    "gpt2-large",
+    "gpt2-xl",
+    "distilgpt2",
+    # See all GPT-2 models at https://huggingface.co/models?filter=gpt2
+]
+
+logger.setLevel(logging.INFO)
+console = logging.StreamHandler()
+console.setLevel(logging.INFO)
+logger.addHandler(console)
+
+_GPT2_ML_TF_TO_TORCH = {
+    "LayerNorm_embed_norm": "emb_norm",
+    "pos_embed": "wpe.weight",
+    "word_embed": "wte.weight",
+    "layer": "h",
+    # Most importently This two layer norm must be put on the same position as gpt2-ml
+    # or generated data is bad, just repeat the last token
+    "LayerNorm_mlp_ln0": "ln_1",
+    "LayerNorm_mlp_ln1": "ln_2",
+    "intermediate": "mlp.c_fc",
+    "output": "mlp.c_proj",
+    "query_layer": "attn.c_attn",
+    "key_layer": "attn.c_attn",
+    "value_layer": "attn.c_attn",
+    "context_projection_layer": "attn.c_proj",
+    "gamma": "weight",
+    "kernel": "weight",
+    "beta": "bias",
+    "bias": "bias",
+}
+
+
+def convert_gpt2_checkpoint_to_pytorch(
+    gpt2_checkpoint_path, gpt2_config_file, pytorch_dump_folder_path
+):
+    # Construct model
+    if gpt2_config_file == "":
+        config = GPT2Config()
+    else:
+        config = GPT2Config.from_json_file(gpt2_config_file)
+    model = GPT2Model(config)
+
+    # Load weights from numpy
+    load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_path)
+
+    # Save pytorch-model
+    pytorch_weights_dump_path = pytorch_dump_folder_path + "/" + WEIGHTS_NAME
+    pytorch_config_dump_path = pytorch_dump_folder_path + "/" + CONFIG_NAME
+    print("Save PyTorch model to {}".format(pytorch_weights_dump_path))
+    torch.save(model.state_dict(), pytorch_weights_dump_path)
+    print("Save configuration file to {}".format(pytorch_config_dump_path))
+    with open(pytorch_config_dump_path, "w", encoding="utf-8") as f:
+        f.write(config.to_json_string())
+
+
+# XXX: MUST do like: convert_gpt2_checkpoint_to_pytorch('./model.ckpt-100000', './mega.json', './')
+#      https://github.com/tensorflow/models/issues/2675#issuecomment-516595597
+def load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_path):
+    """Load tf checkpoints in a pytorch model"""
+    try:
+        import re
+
+        import tensorflow as tf
+    except ImportError:
+        logger.error(
+            "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
+            "https://www.tensorflow.org/install/ for installation instructions."
+        )
+        raise
+    tf_path = os.path.abspath(gpt2_checkpoint_path)
+    logger.info("Converting TensorFlow checkpoint from {}".format(tf_path))
+    # Load weights from TF model
+    init_vars = tf.train.list_variables(tf_path)
+    names = []
+    arrays = []
+    for name, shape in init_vars:
+        logger.info("Loading TF weight {} with shape {}".format(name, shape))
+        array = tf.train.load_variable(tf_path, name)
+        names.append(name)
+        arrays.append(array.squeeze())
+
+    import copy
+
+    orig_model = copy.deepcopy(model)
+
+    for name, array in zip(names, arrays):
+        name = name[6:]  # skip "model/"
+        name = name.split("/")
+        pointer = model
+
+        attn_layer = ""
+        for m_name in name:
+            if re.fullmatch(r"[A-Za-z]+\d+", m_name):
+                scope_names = re.split(r"(\d+)", m_name)
+            else:
+                scope_names = [m_name]
+            sname = scope_names[0]
+
+            if sname == "" or sname == "embeddings":
+                continue
+            elif sname not in _GPT2_ML_TF_TO_TORCH:
+                print("=========================================================")
+                logger.info("Skip var name {}".format(scope_names))
+                pointer = None
+                break
+            else:
+                tname = _GPT2_ML_TF_TO_TORCH[sname]
+                if "." in tname:
+                    parent, child = tname.split(".")
+                    pointer = getattr(pointer, parent)
+                    pointer = getattr(pointer, child)
+                else:
+                    pointer = getattr(pointer, tname)
+
+                if tname == "attn.c_attn":
+                    attn_layer = sname
+
+            if len(scope_names) >= 2:
+                num = int(scope_names[1])
+                pointer = pointer[num]
+
+        if pointer is None:
+            continue
+        if attn_layer == "":
+            try:
+                assert pointer.shape == array.shape
+            except AssertionError as e:
+                e.args += (pointer.shape, array.shape)
+                raise
+        logger.info(
+            "Initialize PyTorch weight {}, {}, {}".format(
+                name, array.mean(), pointer.mean()
+            )
+        )
+        if attn_layer == "":
+            pointer.data = torch.from_numpy(array)
+        else:
+            shape = pointer.shape
+            d = torch.from_numpy(array)
+            is_bias = len(shape) == 1
+            end = int(shape[0 if is_bias else 1] / 3)
+            m = dict(
+                query_layer=0,
+                key_layer=end,
+                value_layer=end * 2,
+            )
+            start = m[attn_layer]
+            end = start + end
+            if is_bias:
+                pointer.data[start:end] = d
+            else:
+                pointer.data[:, start:end] = d
+        logger.info(
+            "Initialize PyTorch weight {}, {}, {}".format(
+                name, array.mean(), pointer.mean()
+            )
+        )
+
+    for name, params in orig_model.named_parameters():
+        for n, p in model.named_parameters():
+            if name == n:
+                if params.equal(p):
+                    print("--------------------------")
+                    print(" %s not changed!" % n)
+    return model
+
+
+class Attention(nn.Module):
+    def __init__(self, nx, n_ctx, config, scale=False, is_cross_attention=False):
+        super().__init__()
+
+        n_state = nx  # in Attention: n_state=768 (nx=n_embd)
+        # [switch nx => n_state from Block to Attention to keep identical to TF implem]
+        assert n_state % config.n_head == 0
+        self.register_buffer(
+            "bias",
+            torch.tril(torch.ones((n_ctx, n_ctx), dtype=torch.uint8)).view(
+                1, 1, n_ctx, n_ctx
+            ),
+        )
+        self.register_buffer("masked_bias", torch.tensor(-1e4))
+        self.n_head = config.n_head
+        self.split_size = n_state
+        self.scale = scale
+        self.is_cross_attention = is_cross_attention
+        if self.is_cross_attention:
+            self.c_attn = Conv1D(2 * n_state, nx)
+            self.q_attn = Conv1D(n_state, nx)
+        else:
+            self.c_attn = Conv1D(3 * n_state, nx)
+        self.c_proj = Conv1D(n_state, nx)
+        self.attn_dropout = nn.Dropout(config.attn_pdrop)
+        self.resid_dropout = nn.Dropout(config.resid_pdrop)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.n_head, self.split_size // self.n_head, self.pruned_heads
+        )
+        index_attn = torch.cat(
+            [index, index + self.split_size, index + (2 * self.split_size)]
+        )
+
+        # Prune conv1d layers
+        self.c_attn = prune_conv1d_layer(self.c_attn, index_attn, dim=1)
+        self.c_proj = prune_conv1d_layer(self.c_proj, index, dim=0)
+
+        # Update hyper params
+        self.split_size = (self.split_size // self.n_head) * (self.n_head - len(heads))
+        self.n_head = self.n_head - len(heads)
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def _attn(
+        self, q, k, v, attention_mask=None, head_mask=None, output_attentions=False
+    ):
+        w = torch.matmul(q, k)
+        if self.scale:
+            w = w / (float(v.size(-1)) ** 0.5)
+        nd, ns = w.size(-2), w.size(-1)
+
+        if not self.is_cross_attention:
+            # if only "normal" attention layer implements causal mask
+            mask = self.bias[:, :, ns - nd : ns, :ns]
+            w = torch.where(mask.bool(), w, self.masked_bias.to(w.dtype))
+
+        if attention_mask is not None:
+            # Apply the attention mask
+            w = w + attention_mask
+
+        w = nn.Softmax(dim=-1)(w)
+        w = self.attn_dropout(w)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            w = w * head_mask
+
+        outputs = [torch.matmul(w, v)]
+        if output_attentions:
+            outputs.append(w)
+        return outputs
+
+    def merge_heads(self, x):
+        x = x.permute(0, 2, 1, 3).contiguous()
+        new_x_shape = x.size()[:-2] + (x.size(-2) * x.size(-1),)
+        return x.view(*new_x_shape)  # in Tensorflow implem: fct merge_states
+
+    def split_heads(self, x, k=False):
+        new_x_shape = x.size()[:-1] + (self.n_head, x.size(-1) // self.n_head)
+        x = x.view(*new_x_shape)  # in Tensorflow implem: fct split_states
+        if k:
+            return x.permute(0, 2, 3, 1)  # (batch, head, head_features, seq_length)
+        else:
+            return x.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)
+
+    def forward(
+        self,
+        hidden_states,
+        layer_past=None,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        use_cache=False,
+        output_attentions=False,
+    ):
+        if encoder_hidden_states is not None:
+            assert hasattr(
+                self, "q_attn"
+            ), "If class is used as cross attention, the weights `q_attn` have to be defined. Please make sure to instantiate class with `Attention(..., is_cross_attention=True)`."
+            query = self.q_attn(hidden_states)
+            key, value = self.c_attn(encoder_hidden_states).split(
+                self.split_size, dim=2
+            )
+            attention_mask = encoder_attention_mask
+        else:
+            query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
+
+        query = self.split_heads(query)
+        key = self.split_heads(key, k=True)
+        value = self.split_heads(value)
+        if layer_past is not None:
+            past_key, past_value = (
+                layer_past[0].transpose(-2, -1),
+                layer_past[1],
+            )  # transpose back cf below
+            key = torch.cat((past_key, key), dim=-1)
+            value = torch.cat((past_value, value), dim=-2)
+
+        if use_cache is True:
+            present = torch.stack(
+                (key.transpose(-2, -1), value)
+            )  # transpose to have same shapes for stacking
+        else:
+            present = (None,)
+
+        attn_outputs = self._attn(
+            query, key, value, attention_mask, head_mask, output_attentions
+        )
+        a = attn_outputs[0]
+
+        a = self.merge_heads(a)
+        a = self.c_proj(a)
+        a = self.resid_dropout(a)
+
+        outputs = [a, present] + attn_outputs[1:]
+        return outputs  # a, present, (attentions)
+
+
+class MLP(nn.Module):
+    def __init__(self, n_state, config):  # in MLP: n_state=3072 (4 * n_embd)
+        super().__init__()
+        nx = config.n_embd
+        self.c_fc = Conv1D(n_state, nx)
+        self.c_proj = Conv1D(nx, n_state)
+        self.act = ACT2FN[config.activation_function]
+        self.dropout = nn.Dropout(config.resid_pdrop)
+
+    def forward(self, x):
+        h = self.act(self.c_fc(x))
+        h2 = self.c_proj(h)
+        return self.dropout(h2)
+
+
+class Block(nn.Module):
+    def __init__(self, n_ctx, config, scale=False):
+        super().__init__()
+        hidden_size = config.n_embd
+        inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
+        self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.attn = Attention(hidden_size, n_ctx, config, scale)
+        self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        if config.add_cross_attention:
+            self.crossattention = Attention(
+                hidden_size, n_ctx, config, scale, is_cross_attention=True
+            )
+            self.ln_cross_attn = nn.LayerNorm(
+                hidden_size, eps=config.layer_norm_epsilon
+            )
+        self.mlp = MLP(inner_dim, config)
+
+    def forward(
+        self,
+        hidden_states,
+        layer_past=None,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        use_cache=False,
+        output_attentions=False,
+    ):
+        attn_outputs = self.attn(
+            hidden_states,
+            layer_past=layer_past,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+        )
+        attn_output = attn_outputs[0]  # output_attn: a, present, (attentions)
+        outputs = attn_outputs[1:]
+        # residual connection
+        hidden_states = attn_output + hidden_states
+
+        if encoder_hidden_states is not None:
+            # add one self-attention block for cross-attention
+            assert hasattr(
+                self, "crossattention"
+            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"
+            cross_attn_outputs = self.crossattention(
+                self.ln_cross_attn(hidden_states),
+                attention_mask=attention_mask,
+                head_mask=head_mask,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
+                output_attentions=output_attentions,
+            )
+            attn_output = cross_attn_outputs[0]
+            # residual connection
+            hidden_states = hidden_states + attn_output
+            outputs = (
+                outputs + cross_attn_outputs[2:]
+            )  # add cross attentions if we output attention weights
+
+        feed_forward_hidden_states = self.mlp(self.ln_1(hidden_states))
+        # residual connection
+        hidden_states = hidden_states + feed_forward_hidden_states
+
+        hidden_states = self.ln_2(hidden_states)
+
+        outputs = [hidden_states] + outputs
+        return outputs  # hidden_states, present, (attentions, cross_attentions)
+
+
+class GPT2PreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = GPT2Config
+    load_tf_weights = load_tf_weights_in_gpt2
+    base_model_prefix = "transformer"
+    is_parallelizable = True
+
+    def __init__(self, *inputs, **kwargs):
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        if isinstance(module, (nn.Linear, nn.Embedding, Conv1D)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if isinstance(module, (nn.Linear, Conv1D)) and module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+@dataclass
+class GPT2DoubleHeadsModelOutput(ModelOutput):
+    """
+    Base class for outputs of models predicting if two sentences are consecutive or not.
+
+    Args:
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided):
+            Language modeling loss.
+        mc_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`mc_labels` is provided):
+            Multiple choice classification loss.
+        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_choices, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        mc_logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_choices)`):
+            Prediction scores of the multiple choice classification head (scores for each choice before SoftMax).
+        past_key_values (:obj:`List[torch.FloatTensor]`, `optional`, returned when ``use_cache=True`` is passed or when ``config.use_cache=True``):
+            List of :obj:`torch.FloatTensor` of length :obj:`config.n_layers`, with each tensor of shape :obj:`(2,
+            batch_size, num_heads, sequence_length, embed_size_per_head)`).
+
+            Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see
+            :obj:`past_key_values` input) to speed up sequential decoding.
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape :obj:`(batch_size, num_heads,
+            sequence_length, sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    mc_loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    mc_logits: torch.FloatTensor = None
+    past_key_values: Optional[List[torch.FloatTensor]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+GPT2_START_DOCSTRING = r"""
+
+    This model inherits from :class:`~transformers.PreTrainedModel`. Check the superclass documentation for the generic
+    methods the library implements for all its model (such as downloading or saving, resizing the input embeddings,
+    pruning heads etc.)
+
+    This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__
+    subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
+    general usage and behavior.
+
+    Parameters:
+        config (:class:`~transformers.GPT2Config`): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model
+            weights.
+"""
+
+GPT2_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, input_ids_length)`):
+            :obj:`input_ids_length` = ``sequence_length`` if :obj:`past_key_values` is ``None`` else
+            ``past_key_values[0].shape[-2]`` (``sequence_length`` of input past key value states). Indices of input
+            sequence tokens in the vocabulary.
+
+            If :obj:`past_key_values` is used, only ``input_ids`` that do not have their past calculated should be
+            passed as ``input_ids``.
+
+            Indices can be obtained using :class:`~transformers.GPT2Tokenizer`. See
+            :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for
+            details.
+
+            `What are input IDs? <../glossary.html#input-ids>`__
+        past_key_values (:obj:`List[torch.FloatTensor]` of length :obj:`config.n_layers`):
+            Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see
+            :obj:`past_key_values` output below). Can be used to speed up sequential decoding. The ``input_ids`` which
+            have their past given to this model should not be passed as ``input_ids`` as they have already been
+            computed.
+        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            `What are attention masks? <../glossary.html#attention-mask>`__
+        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, input_ids_length)`, `optional`):
+            Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0,
+            1]``:
+
+            - 0 corresponds to a `sentence A` token,
+            - 1 corresponds to a `sentence B` token.
+
+            `What are token type IDs? <../glossary.html#token-type-ids>`_
+        position_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
+            config.max_position_embeddings - 1]``.
+
+            `What are position IDs? <../glossary.html#position-ids>`_
+        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`):
+            Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``:
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+
+        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
+            This is useful if you want more control over how to convert :obj:`input_ids` indices into associated
+            vectors than the model's internal embedding lookup matrix.
+
+            If :obj:`past_key_values` is used, optionally only the last :obj:`inputs_embeds` have to be input (see
+            :obj:`past_key_values`).
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        output_attentions (:obj:`bool`, `optional`):
+            Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
+            tensors for more detail.
+        output_hidden_states (:obj:`bool`, `optional`):
+            Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
+            more detail.
+        return_dict (:obj:`bool`, `optional`):
+            Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
+"""
+
+PARALLELIZE_DOCSTRING = r"""
+    This is an experimental feature and is a subject to change at a moment's notice.
+
+    Uses a device map to distribute attention modules of the model across several devices. If no device map is given,
+    it will evenly distribute blocks across all devices.
+
+    Args:
+        device_map (:obj:`Dict[int, list]`, optional, defaults to None):
+            A dictionary that maps attention modules to devices. Note that the embedding module and LMHead are always
+            automatically mapped to the first device (for esoteric reasons). That means that the first device should
+            have fewer attention modules mapped to it than other devices. For reference, the gpt2 models have the
+            following number of attention modules:
+
+                - gpt2: 12
+                - gpt2-medium: 24
+                - gpt2-large: 36
+                - gpt2-xl: 48
+
+    Example::
+
+            # Here is an example of a device map on a machine with 4 GPUs using gpt2-xl, which has a total of 48 attention modules:
+            model = GPT2LMHeadModel.from_pretrained('gpt2-xl')
+            device_map = {0: [0, 1, 2, 3, 4, 5, 6, 7, 8],
+
+                          1: [9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
+                          2: [22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34],
+                          3: [35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47]}
+            model.parallelize(device_map)
+"""
+DEPARALLELIZE_DOCSTRING = r"""
+    Moves the model to cpu from a model parallel state.
+
+    Example::
+
+        # On a 4 GPU machine with gpt2-large:
+        model = GPT2LMHeadModel.from_pretrained('gpt2-large')
+        device_map = {0: [0, 1, 2, 3, 4, 5, 6, 7],
+
+                    1: [8, 9, 10, 11, 12, 13, 14, 15],
+                    2: [16, 17, 18, 19, 20, 21, 22, 23],
+                    3: [24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35]}
+        model.parallelize(device_map) # Splits the model across several devices
+        model.deparallelize() # Put the model back on cpu and cleans memory by calling torch.cuda.empty_cache()
+"""
+
+
+@add_start_docstrings(
+    "The bare GPT2 Model transformer outputting raw hidden-states without any specific head on top.",
+    GPT2_START_DOCSTRING,
+)
+class GPT2Model(GPT2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
+        self.wpe = nn.Embedding(config.n_positions, config.n_embd)
+        if _USE_GROVER:
+            self.emb_norm = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+
+        self.drop = nn.Dropout(config.embd_pdrop)
+        self.h = nn.ModuleList(
+            [Block(config.n_ctx, config, scale=True) for _ in range(config.n_layer)]
+        )
+        if not _USE_GROVER:
+            self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+
+        self.init_weights()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+    @add_start_docstrings(PARALLELIZE_DOCSTRING)
+    def parallelize(self, device_map=None):
+        # Check validity of device_map
+        self.device_map = (
+            get_device_map(len(self.h), range(torch.cuda.device_count()))
+            if device_map is None
+            else device_map
+        )
+        assert_device_map(self.device_map, len(self.h))
+        self.model_parallel = True
+        self.first_device = (
+            "cpu"
+            if "cpu" in self.device_map.keys()
+            else "cuda:" + str(min(self.device_map.keys()))
+        )
+        self.last_device = "cuda:" + str(max(self.device_map.keys()))
+        self.wte = self.wte.to(self.first_device)
+        self.wpe = self.wpe.to(self.first_device)
+        # Load onto devices
+        for k, v in self.device_map.items():
+            for block in v:
+                cuda_device = "cuda:" + str(k)
+                self.h[block] = self.h[block].to(cuda_device)
+        # ln_f to last
+        self.ln_f = self.ln_f.to(self.last_device)
+
+    @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
+    def deparallelize(self):
+        self.model_parallel = False
+        self.device_map = None
+        self.first_device = "cpu"
+        self.last_device = "cpu"
+        self.wte = self.wte.to("cpu")
+        self.wpe = self.wpe.to("cpu")
+        for index in range(len(self.h)):
+            self.h[index] = self.h[index].to("cpu")
+        self.ln_f = self.ln_f.to("cpu")
+        torch.cuda.empty_cache()
+
+    def get_input_embeddings(self):
+        return self.wte
+
+    def set_input_embeddings(self, new_embeddings):
+        self.wte = new_embeddings
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
+        """
+        for layer, heads in heads_to_prune.items():
+            self.h[layer].attn.prune_heads(heads)
+
+    @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint="gpt2",
+        output_type=BaseModelOutputWithPastAndCrossAttentions,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time"
+            )
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+            input_ids = input_ids.view(-1, input_shape[-1])
+            batch_size = input_ids.shape[0]
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+            batch_size = inputs_embeds.shape[0]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if token_type_ids is not None:
+            token_type_ids = token_type_ids.view(-1, input_shape[-1])
+        if position_ids is not None:
+            position_ids = position_ids.view(-1, input_shape[-1])
+
+        if past_key_values is None:
+            past_length = 0
+            past_key_values = [None] * len(self.h)
+        else:
+            past_length = past_key_values[0][0].size(-2)
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_length,
+                input_shape[-1] + past_length,
+                dtype=torch.long,
+                device=device,
+            )
+            position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])
+
+        # Attention mask.
+        if attention_mask is not None:
+            if batch_size <= 0:
+                raise ValueError("batch_size has to be defined and > 0")
+            attention_mask = attention_mask.view(batch_size, -1)
+            # We create a 3D attention mask from a 2D tensor mask.
+            # Sizes are [batch_size, 1, 1, to_seq_length]
+            # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
+            # this attention mask is more simple than the triangular masking of causal attention
+            # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
+            attention_mask = attention_mask[:, None, None, :]
+
+            # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+            # masked positions, this operation will create a tensor which is 0.0 for
+            # positions we want to attend and -10000.0 for masked positions.
+            # Since we are adding it to the raw scores before the softmax, this is
+            # effectively the same as removing these entirely.
+            attention_mask = attention_mask.to(dtype=self.dtype)  # fp16 compatibility
+            attention_mask = (1.0 - attention_mask) * -10000.0
+
+        # If a 2D ou 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if self.config.add_cross_attention and encoder_hidden_states is not None:
+            (
+                encoder_batch_size,
+                encoder_sequence_length,
+                _,
+            ) = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+            encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # head_mask has shape n_layer x batch x n_heads x N x N
+        head_mask = self.get_head_mask(head_mask, self.config.n_layer)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.wte(input_ids)
+        position_embeds = self.wpe(position_ids)
+        hidden_states = inputs_embeds + position_embeds
+
+        if token_type_ids is not None:
+            token_type_embeds = self.wte(token_type_ids)
+            hidden_states = hidden_states + token_type_embeds
+
+        hidden_states = self.drop(hidden_states)
+        if _USE_GROVER:
+            hidden_states = self.emb_norm(hidden_states)
+        output_shape = input_shape + (hidden_states.size(-1),)
+
+        presents = () if use_cache else None
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = (
+            () if output_attentions and self.config.add_cross_attention else None
+        )
+        all_hidden_states = () if output_hidden_states else None
+        for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
+
+            # Model parallel
+            if self.model_parallel:
+                torch.cuda.set_device(hidden_states.device)
+                # Ensure layer_past is on same device as hidden_states (might not be correct)
+                if layer_past is not None:
+                    layer_past = tuple(
+                        past_state.to(hidden_states.device) for past_state in layer_past
+                    )
+                # Ensure that attention_mask is always on the same device as hidden_states
+                if attention_mask is not None:
+                    attention_mask = attention_mask.to(hidden_states.device)
+                if isinstance(head_mask, torch.Tensor):
+                    head_mask = head_mask.to(hidden_states.device)
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (
+                    hidden_states.view(*output_shape),
+                )
+
+            if getattr(self.config, "gradient_checkpointing", False):
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # checkpointing only works with tuple returns, not with lists
+                        return tuple(
+                            output
+                            for output in module(*inputs, use_cache, output_attentions)
+                        )
+
+                    return custom_forward
+
+                outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    layer_past,
+                    attention_mask,
+                    head_mask[i],
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+            else:
+                outputs = block(
+                    hidden_states,
+                    layer_past=layer_past,
+                    attention_mask=attention_mask,
+                    head_mask=head_mask[i],
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                )
+
+            hidden_states, present = outputs[:2]
+            if use_cache is True:
+                presents = presents + (present,)
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (
+                    outputs[2 if use_cache else 1],
+                )
+                if self.config.add_cross_attention:
+                    all_cross_attentions = all_cross_attentions + (
+                        outputs[3 if use_cache else 2],
+                    )
+
+            # Model Parallel: If it's the last layer for that device, put things on the next device
+            if self.model_parallel:
+                for k, v in self.device_map.items():
+                    if i == v[-1] and "cuda:" + str(k) != self.last_device:
+                        hidden_states = hidden_states.to("cuda:" + str(k + 1))
+
+        if not _USE_GROVER:
+            hidden_states = self.ln_f(hidden_states)
+
+        hidden_states = hidden_states.view(*output_shape)
+        # Add last hidden state
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    presents,
+                    all_hidden_states,
+                    all_self_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=presents,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input
+    embeddings).
+    """,
+    GPT2_START_DOCSTRING,
+)
+class GPT2LMHeadModel(GPT2PreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"h\.\d+\.attn\.masked_bias", r"lm_head\.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = GPT2Model(config)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        self.init_weights()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+    @add_start_docstrings(PARALLELIZE_DOCSTRING)
+    def parallelize(self, device_map=None):
+        self.device_map = (
+            get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
+            if device_map is None
+            else device_map
+        )
+        assert_device_map(self.device_map, len(self.transformer.h))
+        self.transformer.parallelize(self.device_map)
+        self.lm_head = self.lm_head.to(self.transformer.first_device)
+        self.model_parallel = True
+
+    @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
+    def deparallelize(self):
+        self.transformer.deparallelize()
+        self.transformer = self.transformer.to("cpu")
+        self.lm_head = self.lm_head.to("cpu")
+        self.model_parallel = False
+        torch.cuda.empty_cache()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)
+        # only last token for inputs_ids if past is defined in kwargs
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint="gpt2",
+        output_type=CausalLMOutputWithCrossAttentions,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+            ``labels = input_ids`` Indices are selected in ``[-100, 0, ..., config.vocab_size]`` All labels set to
+            ``-100`` are ignored (masked), the loss is only computed for labels in ``[0, ..., config.vocab_size]``
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+
+        # Set device for model parallelism
+        if self.model_parallel:
+            torch.cuda.set_device(self.transformer.first_device)
+            hidden_states = hidden_states.to(self.lm_head.weight.device)
+
+        lm_logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
+            )
+
+        if not return_dict:
+            output = (lm_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+        )
+
+    @staticmethod
+    def _reorder_cache(
+        past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
+    ) -> Tuple[Tuple[torch.Tensor]]:
+        """
+        This function is used to re-order the :obj:`past_key_values` cache if
+        :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
+        called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(
+                past_state.index_select(0, beam_idx.to(past_state.device))
+                for past_state in layer_past
+            )
+            for layer_past in past
+        )
+
+
+@add_start_docstrings(
+    """
+The GPT2 Model transformer with a language modeling and a multiple-choice classification head on top e.g. for
+RocStories/SWAG tasks. The two heads are two linear layers. The language modeling head has its weights tied to the
+input embeddings, the classification head takes as input the input of a specified classification token index in the
+input sequence).
+""",
+    GPT2_START_DOCSTRING,
+)
+class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        config.num_labels = 1
+        self.transformer = GPT2Model(config)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.multiple_choice_head = SequenceSummary(config)
+
+        self.init_weights()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+    @add_start_docstrings(PARALLELIZE_DOCSTRING)
+    def parallelize(self, device_map=None):
+        self.device_map = (
+            get_device_map(len(self.transformer.h), range(torch.cuda.device_count()))
+            if device_map is None
+            else device_map
+        )
+        assert_device_map(self.device_map, len(self.transformer.h))
+        self.transformer.parallelize(self.device_map)
+        self.lm_head = self.lm_head.to(self.transformer.first_device)
+        self.multiple_choice_head = self.multiple_choice_head.to(
+            self.transformer.first_device
+        )
+        self.model_parallel = True
+
+    @add_start_docstrings(DEPARALLELIZE_DOCSTRING)
+    def deparallelize(self):
+        self.transformer.deparallelize()
+        self.transformer = self.transformer.to("cpu")
+        self.lm_head = self.lm_head.to("cpu")
+        self.multiple_choice_head = self.multiple_choice_head.to("cpu")
+        self.model_parallel = False
+        torch.cuda.empty_cache()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)
+        # only last token for inputs_ids if past is defined in kwargs
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
+    @replace_return_docstrings(
+        output_type=GPT2DoubleHeadsModelOutput, config_class=_CONFIG_FOR_DOC
+    )
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        mc_token_ids=None,
+        labels=None,
+        mc_labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        **kwargs,
+    ):
+        r"""
+        mc_token_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, num_choices)`, `optional`, default to index of the last token of the input):
+            Index of the classification token in each input sequence. Selected in the range ``[0, input_ids.size(-1) -
+            1[``.
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+            ``labels = input_ids`` Indices are selected in ``[-1, 0, ..., config.vocab_size]`` All labels set to
+            ``-100`` are ignored (masked), the loss is only computed for labels in ``[0, ..., config.vocab_size]``
+        mc_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size)`, `optional`):
+            Labels for computing the multiple choice classification loss. Indices should be in ``[0, ...,
+            num_choices]`` where `num_choices` is the size of the second dimension of the input tensors. (see
+            `input_ids` above)
+
+        Return:
+
+        Example::
+
+            >>> import torch
+            >>> from transformers import GPT2Tokenizer, GPT2DoubleHeadsModel
+
+            >>> tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+            >>> model = GPT2DoubleHeadsModel.from_pretrained('gpt2')
+
+            >>> # Add a [CLS] to the vocabulary (we should train it also!)
+            >>> num_added_tokens = tokenizer.add_special_tokens({'cls_token': '[CLS]'})
+
+            >>> embedding_layer = model.resize_token_embeddings(len(tokenizer))  # Update the model embeddings with the new vocabulary size
+
+            >>> choices = ["Hello, my dog is cute [CLS]", "Hello, my cat is cute [CLS]"]
+            >>> encoded_choices = [tokenizer.encode(s) for s in choices]
+            >>> cls_token_location = [tokens.index(tokenizer.cls_token_id) for tokens in encoded_choices]
+
+            >>> input_ids = torch.tensor(encoded_choices).unsqueeze(0)  # Batch size: 1, number of choices: 2
+            >>> mc_token_ids = torch.tensor([cls_token_location])  # Batch size: 1
+
+            >>> outputs = model(input_ids, mc_token_ids=mc_token_ids)
+            >>> lm_logits = outputs.lm_logits
+            >>> mc_logits = outputs.mc_logits
+
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = transformer_outputs[0]
+
+        # Set device for model parallelism
+        if self.model_parallel:
+            torch.cuda.set_device(self.transformer.first_device)
+            hidden_states = hidden_states.to(self.lm_head.weight.device)
+
+        lm_logits = self.lm_head(hidden_states)
+        mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids).squeeze(-1)
+
+        mc_loss = None
+        if mc_labels is not None:
+            loss_fct = CrossEntropyLoss()
+            mc_loss = loss_fct(
+                mc_logits.view(-1, mc_logits.size(-1)), mc_labels.view(-1)
+            )
+        lm_loss = None
+        if labels is not None:
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = CrossEntropyLoss()
+            lm_loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
+            )
+
+        if not return_dict:
+            output = (lm_logits, mc_logits) + transformer_outputs[1:]
+            if mc_loss is not None:
+                output = (mc_loss,) + output
+            return ((lm_loss,) + output) if lm_loss is not None else output
+
+        return GPT2DoubleHeadsModelOutput(
+            loss=lm_loss,
+            mc_loss=mc_loss,
+            logits=lm_logits,
+            mc_logits=mc_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+    @staticmethod
+    def _reorder_cache(
+        past: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
+    ) -> Tuple[Tuple[torch.Tensor]]:
+        """
+        This function is used to re-order the :obj:`past_key_values` cache if
+        :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
+        called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(
+                past_state.index_select(0, beam_idx.to(past_state.device))
+                for past_state in layer_past
+            )
+            for layer_past in past
+        )
+
+
+@add_start_docstrings(
+    """
+    The GPT2 Model transformer with a sequence classification head on top (linear layer).
+
+    :class:`~transformers.GPT2ForSequenceClassification` uses the last token in order to do the classification, as
+    other causal models (e.g. GPT-1) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    :obj:`pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each
+    row. If no :obj:`pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot
+    guess the padding tokens when :obj:`inputs_embeds` are passed instead of :obj:`input_ids`, it does the same (take
+    the last value in each row of the batch).
+    """,
+    GPT2_START_DOCSTRING,
+)
+class GPT2ForSequenceClassification(GPT2PreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"h\.\d+\.attn\.masked_bias", r"lm_head\.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.transformer = GPT2Model(config)
+        self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
+
+        self.init_weights()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+    @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint="microsoft/dialogrpt",
+        output_type=SequenceClassifierOutputWithPast,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
+            config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size, sequence_length = input_ids.shape[:2]
+        else:
+            batch_size, sequence_length = inputs_embeds.shape[:2]
+
+        assert (
+            self.config.pad_token_id is not None or batch_size == 1
+        ), "Cannot handle batch sizes > 1 if no padding token is defined."
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                sequence_lengths = (
+                    torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1
+                )
+            else:
+                sequence_lengths = -1
+                logger.warning(
+                    f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+                    f"unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+                )
+
+        pooled_logits = logits[range(batch_size), sequence_lengths]
+
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                #  We are doing regression
+                loss_fct = MSELoss()
+                loss = loss_fct(pooled_logits.view(-1), labels.to(self.dtype).view(-1))
+            else:
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(
+                    pooled_logits.view(-1, self.num_labels), labels.view(-1)
+                )
+
+        if not return_dict:
+            output = (pooled_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    GPT2 Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
+    Named-Entity-Recognition (NER) tasks.
+    """,
+    GPT2_START_DOCSTRING,
+)
+class GPT2ForTokenClassification(GPT2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.transformer = GPT2Model(config)
+        if (
+            hasattr(config, "classifier_dropout")
+            and config.classifier_dropout is not None
+        ):
+            classifier_dropout = config.classifier_dropout
+        elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
+            classifier_dropout = config.hidden_dropout
+        else:
+            classifier_dropout = 0.1
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+    @add_start_docstrings_to_model_forward(GPT2_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint="microsoft/DialogRPT-updown",
+        output_type=TokenClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
+            config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = transformer_outputs[0]
+        hidden_states = self.dropout(hidden_states)
+        logits = self.classifier(hidden_states)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            # Only keep active parts of the loss
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)
+                active_labels = torch.where(
+                    active_loss,
+                    labels.view(-1),
+                    torch.tensor(loss_fct.ignore_index).type_as(labels),
+                )
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + transformer_outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )

backend/preprocess.py

@@ -0,0 +1,736 @@
+import html
+import logging
+import re
+from typing import List
+from farasa.segmenter import FarasaSegmenter
+import emoji
+
+import pyarabic.araby as araby
+
+ACCEPTED_MODELS = [
+    "bert-base-arabertv01",
+    "bert-base-arabert",
+    "bert-base-arabertv02",
+    "bert-base-arabertv2",
+    "bert-large-arabertv02",
+    "bert-large-arabertv2",
+    "araelectra-base",
+    "araelectra-base-discriminator",
+    "araelectra-base-generator",
+    "araelectra-base-artydiqa",
+    "aragpt2-base",
+    "aragpt2-medium",
+    "aragpt2-large",
+    "aragpt2-mega",
+]
+
+SEGMENTED_MODELS = [
+    "bert-base-arabert",
+    "bert-base-arabertv2",
+    "bert-large-arabertv2",
+]
+
+SECOND_GEN_MODELS = [
+    "bert-base-arabertv02",
+    "bert-base-arabertv2",
+    "bert-large-arabertv02",
+    "bert-large-arabertv2",
+    "araelectra-base",
+    "araelectra-base-discriminator",
+    "araelectra-base-generator",
+    "araelectra-base-artydiqa",
+    "aragpt2-base",
+    "aragpt2-medium",
+    "aragpt2-large",
+    "aragpt2-mega",
+]
+
+farasa_segmenter = FarasaSegmenter(interactive=True)
+
+
+class ArabertPreprocessor:
+    """
+    A Preprocessor class that cleans and preprocesses text for all models in the AraBERT repo.
+    It also can unprocess the text ouput of the generated text
+
+    Args:
+
+        model_name (:obj:`str`): model name from the HuggingFace Models page without
+        the aubmindlab tag. Will default to a base Arabic preprocessor if model name was not found.
+        Current accepted models are:
+
+            - "bert-base-arabertv01": No farasa segmentation.
+            - "bert-base-arabert": with farasa segmentation.
+            - "bert-base-arabertv02": No farasas egmentation.
+            - "bert-base-arabertv2": with farasa segmentation.
+            - "bert-large-arabertv02": No farasas egmentation.
+            - "bert-large-arabertv2": with farasa segmentation.
+            - "araelectra-base": No farasa segmentation.
+            - "araelectra-base-discriminator": No farasa segmentation.
+            - "araelectra-base-generator": No farasa segmentation.
+            - "aragpt2-base": No farasa segmentation.
+            - "aragpt2-medium": No farasa segmentation.
+            - "aragpt2-large": No farasa segmentation.
+            - "aragpt2-mega": No farasa segmentation.
+
+
+        keep_emojis(:obj:`bool`, `optional`, defaults to :obj:`False`): don't remove emojis while preprocessing.
+
+        remove_html_markup(:obj: `bool`, `optional`, defaults to :obj:`True`): Whether to remove html artfacts,
+        should be set to False when preprocessing TyDi QA.
+
+        replace_urls_emails_mentions(:obj:`bool`, `optional`, defaults to :obj:`True`): Whether to replace email urls
+        and mentions by special tokens.
+
+        strip_tashkeel(:obj:`bool`, `optional`, defaults to :obj:`True`): remove diacritics (FATHATAN, DAMMATAN, KASRATAN, FATHA, DAMMA,
+        KASRA, SUKUN, SHADDA).
+
+        strip_tatweel(:obj:`bool`, `optional`, defaults to :obj:`True`): remove tatweel '\\u0640'.
+
+        insert_white_spaces(:obj:`bool`, `optional`, defaults to :obj:`True`): insert whitespace before and after all non Arabic digits
+        or English digits or Arabic and English Alphabet or the 2 brackets, then inserts whitespace
+        between words and numbers or numbers and words.
+
+        remove_non_digit_repetition(:obj:`bool`, `optional`, defaults to :obj:`True`): replace repetition of more than 2 non-digit character with
+        2 of this character.
+
+        replace_slash_with_dash(:obj:`bool`, `optional`, defaults to :obj:`None`): Will be automatically set to True in AraBERTv02,
+        AraELECTRA and AraGPT2.
+        Set to False to force disable, and True to force enable. Replaces the "/"  with "-",
+        since "/" is missing from AraBERTv2, AraELECTRA and ARAGPT2 vocabulary.
+
+        map_hindi_numbers_to_arabic(:obj:`bool`, `optional`, defaults to :obj:`None`): Will be automatically set to True in
+        AraBERTv02, AraELECTRA and AraGPT2.Set to False to force disable, and True to force enable.
+        Replaces hindi numbers with the corresponding Arabic one. ex: "١٩٩٥" --> "1995".
+        This is behavior is present by default in AraBERTv1 and v2 (with pre-segmentation),
+        and fixes the issue of caused by a bug when inserting white spaces.
+
+        apply_farasa_segmentation(:obj:`bool`, `optional`, defaults to :obj:`None`): Will be automatically set to True in
+        AraBERTv2, and AraBERTv1. Set to False to force disable, and True to force enable.
+
+
+
+    Returns:
+
+        ArabertPreprocessor: A preprocessor instance
+
+    Example:
+
+        from preprocess import ArabertPreprocessor
+
+        arabert_prep = ArabertPreprocessor("aubmindlab/bert-base-arabertv2")
+
+        arabert_prep.preprocess("SOME ARABIC TEXT")
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        keep_emojis: bool = False,
+        remove_html_markup: bool = True,
+        replace_urls_emails_mentions: bool = True,
+        strip_tashkeel: bool = True,
+        strip_tatweel: bool = True,
+        insert_white_spaces: bool = True,
+        remove_non_digit_repetition: bool = True,
+        replace_slash_with_dash: bool = None,
+        map_hindi_numbers_to_arabic: bool = None,
+        apply_farasa_segmentation: bool = None,
+    ):
+
+        model_name = model_name.replace("aubmindlab/", "").replace("wissamantoun/", "")
+
+        if model_name not in ACCEPTED_MODELS:
+            logging.warning(
+                """Model provided is not in the accepted model list. Preprocessor will default to a base Arabic preprocessor"""
+            )
+            self.model_name = "bert-base-arabertv02"
+        else:
+            self.model_name = model_name
+
+        if apply_farasa_segmentation is None:
+            if self.model_name in SEGMENTED_MODELS:
+                self.apply_farasa_segmentation = True
+            else:
+                self.apply_farasa_segmentation = False
+        else:
+            if (
+                apply_farasa_segmentation == False
+                and self.model_name in SEGMENTED_MODELS
+            ):
+                logging.warning(
+                    "The selected model_name requires Farasa pre-segmentation, but apply_farasa_segmentation was set to False!"
+                )
+
+            self.apply_farasa_segmentation = apply_farasa_segmentation
+
+        self.keep_emojis = keep_emojis
+        self.remove_html_markup = remove_html_markup
+        self.replace_urls_emails_mentions = replace_urls_emails_mentions
+        self.strip_tashkeel = strip_tashkeel
+        self.strip_tatweel = strip_tatweel
+        self.insert_white_spaces = insert_white_spaces
+        self.remove_non_digit_repetition = remove_non_digit_repetition
+
+        if replace_slash_with_dash is None:
+            if self.model_name in SECOND_GEN_MODELS:
+                self.replace_slash_with_dash = True
+            else:
+                self.replace_slash_with_dash = False
+        else:
+            self.replace_slash_with_dash = replace_slash_with_dash
+
+        if map_hindi_numbers_to_arabic is None:
+            if self.model_name in SECOND_GEN_MODELS:
+                self.map_hindi_numbers_to_arabic = True
+            else:
+                self.map_hindi_numbers_to_arabic = False
+        else:
+            self.map_hindi_numbers_to_arabic = map_hindi_numbers_to_arabic
+
+    def preprocess(self, text: str) -> str:
+        """
+        Preprocess takes an input text line an applies the same preprocessing used in AraBERT
+                            pretraining, or according to settings
+
+        Args:
+
+            text (:obj:`str`): inout text string
+
+        Returns:
+
+            string: A preprocessed string depending on which model was selected
+        """
+        if (
+            self.model_name == "bert-base-arabert"
+            or self.model_name == "bert-base-arabertv01"
+        ):
+            return self._preprocess_v1(
+                text,
+                do_farasa_tokenization=self.apply_farasa_segmentation,
+            )
+
+        if self.model_name in SECOND_GEN_MODELS:
+            return self._preprocess_v2(text)
+
+        return self._preprocess_v3(text)
+
+    def unpreprocess(self, text: str, desegment: bool = True) -> str:
+        """Re-formats the text to a classic format where punctuations, brackets, parenthesis are not seperated by whitespaces.
+        The objective is to make the generated text of any model appear natural and not preprocessed.
+
+        Args:
+            text (:obj:`str`): input text to be un-preprocessed
+            desegment (:obj:`bool`, optional): [whether or not to remove farasa pre-segmentation before]..
+
+        Returns:
+            str: The unpreprocessed (and possibly Farasa-desegmented) text.
+        """
+
+        if self.apply_farasa_segmentation and desegment:
+            text = self.desegment(text)
+
+        # removes the spaces around quotation marks ex: i " ate " an apple --> i "ate" an apple
+        # https://stackoverflow.com/a/53436792/5381220
+        text = re.sub(white_spaced_double_quotation_regex, '"' + r"\1" + '"', text)
+        text = re.sub(white_spaced_single_quotation_regex, "'" + r"\1" + "'", text)
+        text = re.sub(white_spaced_back_quotation_regex, "\`" + r"\1" + "\`", text)
+        text = re.sub(white_spaced_back_quotation_regex, "\—" + r"\1" + "\—", text)
+
+        # during generation, sometimes the models don't put a space after the dot, this handles it
+        text = text.replace(".", " . ")
+        text = " ".join(text.split())
+
+        # handle decimals
+        text = re.sub(r"(\d+) \. (\d+)", r"\1.\2", text)
+        text = re.sub(r"(\d+) \, (\d+)", r"\1,\2", text)
+
+        text = re.sub(left_and_right_spaced_chars, r"\1", text)
+        text = re.sub(left_spaced_chars, r"\1", text)
+        text = re.sub(right_spaced_chars, r"\1", text)
+
+        return text
+
+    def desegment(self, text: str) -> str:
+        """
+        Use this function if sentence tokenization was done using
+        `from arabert.preprocess_arabert import preprocess` with Farasa enabled
+        AraBERT segmentation using Farasa adds a space after the '+' for prefixes,
+        and after before the '+' for suffixes
+
+        Example:
+        >>> desegment('ال+ دراس +ات')
+        الدراسات
+        """
+        text = text.replace("+ ", "+")
+        text = text.replace(" +", "+")
+        text = " ".join([self._desegmentword(word) for word in text.split(" ")])
+        return text
+
+    def _desegmentword(self, orig_word: str) -> str:
+        """
+        Word segmentor that takes a Farasa Segmented Word and removes the '+' signs
+
+        Example:
+        >>> _desegmentword("ال+يومي+ة")
+        اليومية
+        """
+        word = orig_word.replace("ل+ال+", "لل")
+        if "ال+ال" not in orig_word:
+            word = word.replace("ل+ال", "لل")
+        word = word.replace("+", "")
+        word = word.replace("للل", "لل")
+        return word
+
+    def _preprocess_v3(self, text: str) -> str:
+        text = str(text)
+        text = html.unescape(text)
+        if self.strip_tashkeel:
+            text = araby.strip_tashkeel(text)
+        if self.strip_tatweel:
+            text = araby.strip_tatweel(text)
+
+        if self.replace_urls_emails_mentions:
+            # replace all possible URLs
+            for reg in url_regexes:
+                text = re.sub(reg, " [رابط] ", text)
+            # REplace Emails with [بريد]
+            for reg in email_regexes:
+                text = re.sub(reg, " [بريد] ", text)
+            # replace mentions with [مستخدم]
+            text = re.sub(user_mention_regex, " [مستخدم] ", text)
+
+        if self.remove_html_markup:
+            # remove html line breaks
+            text = re.sub("<br />", " ", text)
+            # remove html markup
+            text = re.sub("</?[^>]+>", " ", text)
+
+        if self.map_hindi_numbers_to_arabic:
+            text = text.translate(hindi_to_arabic_map)
+
+        # remove repeated characters >2
+        if self.remove_non_digit_repetition:
+            text = self._remove_non_digit_repetition(text)
+
+        # insert whitespace before and after all non Arabic digits or English Digits and Alphabet and the 2 brackets
+        if self.insert_white_spaces:
+            text = re.sub(
+                "([^0-9\u0621-\u063A\u0641-\u064A\u0660-\u0669a-zA-Z ])",
+                r" \1 ",
+                text,
+            )
+
+            # re-fix brackets
+            text = text.replace("[ رابط ]", "[رابط]")
+            text = text.replace("[ بريد ]", "[بريد]")
+            text = text.replace("[ مستخدم ]", "[مستخدم]")
+
+            # insert whitespace between words and numbers or numbers and words
+            text = re.sub(
+                "(\d+)([\u0621-\u063A\u0641-\u064A\u066A-\u066C\u0654-\u0655]+)",
+                r" \1 \2 ",
+                text,
+            )
+            text = re.sub(
+                "([\u0621-\u063A\u0641-\u064A\u066A-\u066C\u0654-\u0655]+)(\d+)",
+                r" \1 \2 ",
+                text,
+            )
+
+        # remove unwanted characters
+        if self.keep_emojis:
+            emoji_regex = "".join(list(emoji.UNICODE_EMOJI["en"].keys()))
+            rejected_chars_regex2 = "[^%s%s]" % (chars_regexv2, emoji_regex)
+            text = re.sub(rejected_chars_regex2, " ", text)
+        else:
+            text = re.sub(rejected_chars_regexv2, " ", text)
+
+        # remove extra spaces
+        text = " ".join(text.replace("\uFE0F", "").split())
+
+        if self.apply_farasa_segmentation:
+            if self.keep_emojis:
+                new_text = []
+                for word in text.split():
+                    if word in list(emoji.UNICODE_EMOJI["en"].keys()):
+                        new_text.append(word)
+                    else:
+                        new_text.append(farasa_segmenter.segment(word))
+                text = " ".join(new_text)
+            else:
+                text = farasa_segmenter.segment(text)
+            return self._farasa_segment(text)
+
+        # ALl the other models dont require Farasa Segmentation
+        return text
+
+    def _preprocess_v2(self, text: str) -> str:
+        text = str(text)
+        text = html.unescape(text)
+        if self.strip_tashkeel:
+            text = araby.strip_tashkeel(text)
+        if self.strip_tatweel:
+            text = araby.strip_tatweel(text)
+
+        if self.replace_urls_emails_mentions:
+            # replace all possible URLs
+            for reg in url_regexes:
+                text = re.sub(reg, " [رابط] ", text)
+            # REplace Emails with [بريد]
+            for reg in email_regexes:
+                text = re.sub(reg, " [بريد] ", text)
+            # replace mentions with [مستخدم]
+            text = re.sub(user_mention_regex, " [مستخدم] ", text)
+
+        if self.remove_html_markup:
+            # remove html line breaks
+            text = re.sub("<br />", " ", text)
+            # remove html markup
+            text = re.sub("</?[^>]+>", " ", text)
+
+        if self.map_hindi_numbers_to_arabic:
+            text = text.translate(hindi_to_arabic_map)
+
+        # remove repeated characters >2
+        if self.remove_non_digit_repetition:
+            text = self._remove_non_digit_repetition(text)
+
+        # insert whitespace before and after all non Arabic digits or English Digits and Alphabet and the 2 brackets
+        if self.insert_white_spaces:
+            text = re.sub(
+                "([^0-9\u0621-\u063A\u0641-\u064A\u0660-\u0669a-zA-Z\[\]])",
+                r" \1 ",
+                text,
+            )
+
+            # insert whitespace between words and numbers or numbers and words
+            text = re.sub(
+                "(\d+)([\u0621-\u063A\u0641-\u064A\u0660-\u066C]+)", r" \1 \2 ", text
+            )
+            text = re.sub(
+                "([\u0621-\u063A\u0641-\u064A\u0660-\u066C]+)(\d+)", r" \1 \2 ", text
+            )
+
+        if self.replace_slash_with_dash:
+            text = text.replace("/", "-")
+
+        # remove unwanted characters
+        if self.keep_emojis:
+            emoji_regex = "".join(list(emoji.UNICODE_EMOJI["en"].keys()))
+            rejected_chars_regex2 = "[^%s%s]" % (chars_regex, emoji_regex)
+            text = re.sub(rejected_chars_regex2, " ", text)
+        else:
+            text = re.sub(rejected_chars_regex, " ", text)
+
+        # remove extra spaces
+        text = " ".join(text.replace("\uFE0F", "").split())
+
+        if (
+            self.model_name == "bert-base-arabertv2"
+            or self.model_name == "bert-large-arabertv2"
+        ):
+            if self.keep_emojis:
+                new_text = []
+                for word in text.split():
+                    if word in list(emoji.UNICODE_EMOJI["en"].keys()):
+                        new_text.append(word)
+                    else:
+                        new_text.append(farasa_segmenter.segment(word))
+                text = " ".join(new_text)
+            else:
+                text = farasa_segmenter.segment(text)
+            return self._farasa_segment(text)
+
+        # ALl the other models dont require Farasa Segmentation
+        return text
+
+    def _preprocess_v1(self, text: str, do_farasa_tokenization: bool) -> str:
+        """
+        AraBERTv1 preprocessing Function
+        """
+        text = str(text)
+        if self.strip_tashkeel:
+            text = araby.strip_tashkeel(text)
+
+        text = re.sub(r"\d+\/[ء-ي]+\/\d+\]", "", text)
+        text = re.sub("ـ", "", text)
+        text = re.sub("[«»]", ' " ', text)
+
+        if self.replace_urls_emails_mentions:
+            # replace the [رابط] token with space if you want to clean links
+            text = re.sub(regex_url_step1, "[رابط]", text)
+            text = re.sub(regex_url_step2, "[رابط]", text)
+            text = re.sub(regex_url, "[رابط]", text)
+            text = re.sub(regex_email, "[بريد]", text)
+            text = re.sub(regex_mention, "[مستخدم]", text)
+        text = re.sub("…", r"\.", text).strip()
+        text = self._remove_redundant_punct(text)
+
+        if self.replace_urls_emails_mentions:
+            text = re.sub(r"\[ رابط \]|\[ رابط\]|\[رابط \]", " [رابط] ", text)
+            text = re.sub(r"\[ بريد \]|\[ بريد\]|\[بريد \]", " [بريد] ", text)
+            text = re.sub(r"\[ مستخدم \]|\[ مستخدم\]|\[مستخدم \]", " [مستخدم] ", text)
+
+        if self.remove_non_digit_repetition:
+            text = self._remove_non_digit_repetition(text)
+
+        if self.insert_white_spaces:
+            text = re.sub(
+                "([^0-9\u0621-\u063A\u0641-\u0669\u0671-\u0673a-zA-Z\[\]])",
+                r" \1 ",
+                text,
+            )
+        if do_farasa_tokenization:
+            text = self._tokenize_arabic_words_farasa(text)
+
+        text = " ".join(text.split())
+
+        return text
+
+    def _farasa_segment(self, text: str) -> str:
+        line_farasa = text.split()
+        segmented_line = []
+        for index, word in enumerate(line_farasa):
+            if word in ["[", "]"]:
+                continue
+            if word in ["رابط", "بريد", "مستخدم"] and line_farasa[index - 1] in [
+                "[",
+                "]",
+            ]:
+                segmented_line.append("[" + word + "]")
+                continue
+            if "+" not in word:
+                segmented_line.append(word)
+                continue
+            segmented_word = self._split_farasa_output(word)
+            segmented_line.extend(segmented_word)
+
+        return " ".join(segmented_line)
+
+    def _split_farasa_output(self, word: str) -> str:
+        segmented_word = []
+        temp_token = ""
+        for i, c in enumerate(word):
+            if c == "+":
+                # if the token is KAF, it could be a suffix or prefix
+                if temp_token == "ك":
+                    # if we are at the second token, then KAF is surely a prefix
+                    if i == 1:
+                        segmented_word.append(temp_token + "+")
+                        temp_token = ""
+                    # If the KAF token is between 2 tokens
+                    elif word[i - 2] == "+":
+                        # if the previous token is prefix, then this KAF must be a prefix
+                        if segmented_word[-1][-1] == "+":
+                            segmented_word.append(temp_token + "+")
+                            temp_token = ""
+                        # else it is a suffix, this KAF could not be a second suffix
+                        else:
+                            segmented_word.append("+" + temp_token)
+                            temp_token = ""
+                    # if Kaf is at the end, this is handled with the statement after the loop
+                elif temp_token in prefix_list:
+                    segmented_word.append(temp_token + "+")
+                    temp_token = ""
+                elif temp_token in suffix_list:
+                    segmented_word.append("+" + temp_token)
+                    temp_token = ""
+                else:
+                    segmented_word.append(temp_token)
+                    temp_token = ""
+                continue
+            temp_token += c
+        if temp_token != "":
+            if temp_token in suffix_list:
+                segmented_word.append("+" + temp_token)
+            else:
+                segmented_word.append(temp_token)
+        return segmented_word
+
+    def _tokenize_arabic_words_farasa(self, line_input: str) -> str:
+
+        if self.keep_emojis:
+            # insert whitespace before and after all non Arabic digits or English Digits and Alphabet and the 2 brackets
+            line_farasa = []
+            for word in line_input.split():
+                if word in list(emoji.UNICODE_EMOJI["en"].keys()):
+                    line_farasa.append(word)
+                else:
+                    line_farasa.append(farasa_segmenter.segment(word))
+        else:
+            line_farasa = farasa_segmenter.segment(line_input).split()
+
+        segmented_line = []
+        for index, word in enumerate(line_farasa):
+            if word in ["[", "]"]:
+                continue
+            if word in ["رابط", "بريد", "مستخدم"] and line_farasa[index - 1] in [
+                "[",
+                "]",
+            ]:
+                segmented_line.append("[" + word + "]")
+                continue
+            segmented_word = []
+            for token in word.split("+"):
+                if token in prefix_list:
+                    segmented_word.append(token + "+")
+                elif token in suffix_list:
+                    segmented_word.append("+" + token)
+                else:
+                    segmented_word.append(token)
+            segmented_line.extend(segmented_word)
+        return " ".join(segmented_line)
+
+    def _remove_non_digit_repetition(self, text: str) -> str:
+        """
+        :param text:  the input text to remove elongation
+        :return: delongated text
+        """
+        # loop over the number of times the regex matched the text
+        # OLD
+        # for index_ in range(len(re.findall(regex_tatweel, text))):
+        #     elongation = re.search(regex_tatweel, text)
+        #     if elongation:
+        #         elongation_pattern = elongation.group()
+        #         elongation_replacement = elongation_pattern[0]
+        #         elongation_pattern = re.escape(elongation_pattern)
+        #         text = re.sub(
+        #             elongation_pattern, elongation_replacement, text, flags=re.MULTILINE
+        #         )
+        #     else:
+        #         break
+
+        # New
+        text = multiple_char_pattern.sub(r"\1\1", text)
+        return text
+
+    def _remove_redundant_punct(self, text: str) -> str:
+        text_ = text
+        result = re.search(redundant_punct_pattern, text)
+        dif = 0
+        while result:
+            sub = result.group()
+            sub = sorted(set(sub), key=sub.index)
+            sub = " " + "".join(list(sub)) + " "
+            text = "".join(
+                (text[: result.span()[0] + dif], sub, text[result.span()[1] + dif :])
+            )
+            text_ = "".join(
+                (text_[: result.span()[0]], text_[result.span()[1] :])
+            ).strip()
+            dif = abs(len(text) - len(text_))
+            result = re.search(redundant_punct_pattern, text_)
+        text = re.sub(r"\s+", " ", text)
+        return text.strip()
+
+
+prefix_list = [
+    "ال",
+    "و",
+    "ف",
+    "ب",
+    "ك",
+    "ل",
+    "لل",
+    "\u0627\u0644",
+    "\u0648",
+    "\u0641",
+    "\u0628",
+    "\u0643",
+    "\u0644",
+    "\u0644\u0644",
+    "س",
+]
+suffix_list = [
+    "ه",
+    "ها",
+    "ك",
+    "ي",
+    "هما",
+    "كما",
+    "نا",
+    "كم",
+    "هم",
+    "هن",
+    "كن",
+    "ا",
+    "ان",
+    "ين",
+    "ون",
+    "وا",
+    "ات",
+    "ت",
+    "ن",
+    "ة",
+    "\u0647",
+    "\u0647\u0627",
+    "\u0643",
+    "\u064a",
+    "\u0647\u0645\u0627",
+    "\u0643\u0645\u0627",
+    "\u0646\u0627",
+    "\u0643\u0645",
+    "\u0647\u0645",
+    "\u0647\u0646",
+    "\u0643\u0646",
+    "\u0627",
+    "\u0627\u0646",
+    "\u064a\u0646",
+    "\u0648\u0646",
+    "\u0648\u0627",
+    "\u0627\u062a",
+    "\u062a",
+    "\u0646",
+    "\u0629",
+]
+other_tokens = ["[رابط]", "[مستخدم]", "[بريد]"]
+
+# the never_split list is ussed with the transformers library
+prefix_symbols = [x + "+" for x in prefix_list]
+suffix_symblos = ["+" + x for x in suffix_list]
+never_split_tokens = list(set(prefix_symbols + suffix_symblos + other_tokens))
+
+url_regexes = [
+    r"(http(s)?:\/\/.)?(www\.)?[-a-zA-Z0-9@:%._\+~#=]{2,256}\.[a-z]{2,6}\b([-a-zA-Z0-9@:%_\+.~#?&//=]*)",
+    r"@(https?|ftp)://(-\.)?([^\s/?\.#-]+\.?)+(/[^\s]*)?$@iS",
+    r"http[s]?://[a-zA-Z0-9_\-./~\?=%&]+",
+    r"www[a-zA-Z0-9_\-?=%&/.~]+",
+    r"[a-zA-Z]+\.com",
+    r"(?=http)[^\s]+",
+    r"(?=www)[^\s]+",
+    r"://",
+]
+user_mention_regex = r"@[\w\d]+"
+email_regexes = [r"[\w-]+@([\w-]+\.)+[\w-]+", r"\S+@\S+"]
+redundant_punct_pattern = (
+    r"([!\"#\$%\'\(\)\*\+,\.:;\-<=·>?@\[\\\]\^_ـ`{\|}~—٪’،؟`୍“؛”ۚ【»؛\s+«–…‘]{2,})"
+)
+
+regex_tatweel = r"(\D)\1{2,}"
+multiple_char_pattern = re.compile(r"(\D)\1{2,}", re.DOTALL)
+
+rejected_chars_regex = r"[^0-9\u0621-\u063A\u0640-\u066C\u0671-\u0674a-zA-Z\[\]!\"#\$%\'\(\)\*\+,\.:;\-<=·>?@\[\\\]\^_ـ`{\|}~—٪’،؟`୍“؛”ۚ»؛\s+«–…‘]"
+rejected_chars_regexv2 = r"[^0-9\u0621-\u063A\u0641-\u066C\u0671-\u0674a-zA-Z\[\]!\"#\$%\'\(\)\*\+,\.:;\-<=·>?@\[\\\]\^_ـ`{\|}~—٪’،؟`୍“؛”ۚ»؛\s+«–…‘/]"
+
+regex_url_step1 = r"(?=http)[^\s]+"
+regex_url_step2 = r"(?=www)[^\s]+"
+regex_url = r"(http(s)?:\/\/.)?(www\.)?[-a-zA-Z0-9@:%._\+~#=]{2,256}\.[a-z]{2,6}\b([-a-zA-Z0-9@:%_\+.~#?&//=]*)"
+regex_mention = r"@[\w\d]+"
+regex_email = r"\S+@\S+"
+
+chars_regex = r"0-9\u0621-\u063A\u0640-\u066C\u0671-\u0674a-zA-Z\[\]!\"#\$%\'\(\)\*\+,\.:;\-<=·>?@\[\\\]\^_ـ`{\|}~—٪’،؟`୍“؛”ۚ»؛\s+«–…‘"
+chars_regexv2 = r"0-9\u0621-\u063A\u0640-\u066C\u0671-\u0674a-zA-Z\[\]!\"#\$%\'\(\)\*\+,\.:;\-<=·>?@\[\\\]\^_ـ`{\|}~—٪’،؟`୍“؛”ۚ»؛\s+«–…‘/"
+
+white_spaced_double_quotation_regex = r'\"\s+([^"]+)\s+\"'
+white_spaced_single_quotation_regex = r"\'\s+([^']+)\s+\'"
+white_spaced_back_quotation_regex = r"\`\s+([^`]+)\s+\`"
+white_spaced_em_dash = r"\—\s+([^—]+)\s+\—"
+
+left_spaced_chars = r" ([\]!#\$%\),\.:;\?}٪’،؟”؛…»·])"
+right_spaced_chars = r"([\[\(\{“«‘*\~]) "
+left_and_right_spaced_chars = r" ([\+\-\<\=\>\@\\\^\_\|\–]) "
+
+hindi_nums = "٠١٢٣٤٥٦٧٨٩"
+arabic_nums = "0123456789"
+hindi_to_arabic_map = str.maketrans(hindi_nums, arabic_nums)

backend/processor.py

@@ -0,0 +1,183 @@
+import streamlit as st
+import awesome_streamlit as ast
+from .preprocess import (
+    ArabertPreprocessor,
+    white_spaced_back_quotation_regex,
+    white_spaced_double_quotation_regex,
+    white_spaced_em_dash,
+    white_spaced_single_quotation_regex,
+    left_and_right_spaced_chars,
+    left_spaced_chars,
+    right_spaced_chars,
+)
+import re
+
+MODELS_to_SELECT = [
+    "None",
+    "bert-base-arabertv01",
+    "bert-base-arabert",
+    "bert-base-arabertv02",
+    "bert-base-arabertv2",
+    "bert-large-arabertv02",
+    "bert-large-arabertv2",
+    "araelectra-base",
+    "araelectra-base-discriminator",
+    "araelectra-base-generator",
+    "araelectra-base-artydiqa",
+    "aragpt2-base",
+    "aragpt2-medium",
+    "aragpt2-large",
+    "aragpt2-mega",
+]
+
+
+def unpreprocess(text: str) -> str:
+    """Re-formats the text to a classic format where punctuations, brackets, parenthesis are not seperated by whitespaces.
+    The objective is to make the generated text of any model appear natural and not preprocessed.
+
+    Args:
+        text (:obj:`str`): input text to be un-preprocessed
+        desegment (:obj:`bool`, optional): [whether or not to remove farasa pre-segmentation before]..
+
+    Returns:
+        str: The unpreprocessed (and possibly Farasa-desegmented) text.
+    """
+
+    text = desegment(text)
+
+    # removes the spaces around quotation marks ex: i " ate " an apple --> i "ate" an apple
+    # https://stackoverflow.com/a/53436792/5381220
+    text = re.sub(white_spaced_double_quotation_regex, '"' + r"\1" + '"', text)
+    text = re.sub(white_spaced_single_quotation_regex, "'" + r"\1" + "'", text)
+    text = re.sub(white_spaced_back_quotation_regex, "\`" + r"\1" + "\`", text)
+    text = re.sub(white_spaced_back_quotation_regex, "\—" + r"\1" + "\—", text)
+
+    # during generation, sometimes the models don't put a space after the dot, this handles it
+    text = text.replace(".", " . ")
+    text = " ".join(text.split())
+
+    # handle decimals
+    text = re.sub(r"(\d+) \. (\d+)", r"\1.\2", text)
+    text = re.sub(r"(\d+) \, (\d+)", r"\1,\2", text)
+
+    text = re.sub(left_and_right_spaced_chars, r"\1", text)
+    text = re.sub(left_spaced_chars, r"\1", text)
+    text = re.sub(right_spaced_chars, r"\1", text)
+
+    return text
+
+
+def desegment(text: str) -> str:
+    """
+    Use this function if sentence tokenization was done using
+    `from arabert.preprocess_arabert import preprocess` with Farasa enabled
+    AraBERT segmentation using Farasa adds a space after the '+' for prefixes,
+    and after before the '+' for suffixes
+
+    Example:
+    >>> desegment('ال+ دراس +ات')
+    الدراسات
+    """
+    text = text.replace("+ ", "+")
+    text = text.replace(" +", "+")
+    text = " ".join([_desegmentword(word) for word in text.split(" ")])
+    return text
+
+
+def _desegmentword(orig_word: str) -> str:
+    """
+    Word segmentor that takes a Farasa Segmented Word and removes the '+' signs
+
+    Example:
+    >>> _desegmentword("ال+يومي+ة")
+    اليومية
+    """
+    word = orig_word.replace("ل+ال+", "لل")
+    if "ال+ال" not in orig_word:
+        word = word.replace("ل+ال", "لل")
+    word = word.replace("+", "")
+    word = word.replace("للل", "لل")
+    return word
+
+
+def write():
+
+    st.markdown(
+        """
+        <h1 style="text-align:left;">Arabic Text Pre-Processor</h1>
+        """,
+        unsafe_allow_html=True,
+    )
+    st.markdown(
+        """
+        <style>
+        p, div, input, label {
+        text-align: right;
+        }
+        </style>
+        """,
+        unsafe_allow_html=True,
+    )
+    input_text = st.text_input(
+        "Text to Pre-Process",
+        value="ولن نبالغ إذا قلنا: إن 'هاتف' أو 'كمبيوتر المكتب' في زمننا هذا ضروري",
+    )
+
+    st.sidebar.title("Model Selector")
+    model_selector = st.sidebar.selectbox(
+        """Select None to enable further filters""", options=MODELS_to_SELECT, index=3
+    )
+    if model_selector == "None":
+        keep_emojis = st.sidebar.checkbox("Keep emojis", False)
+        remove_html_markup = st.sidebar.checkbox("Remove html markup", True)
+        strip_tashkeel = st.sidebar.checkbox("Strip tashkeel", True)
+        replace_urls_emails_mentions = st.sidebar.checkbox(
+            "Replace urls and emails", True
+        )
+        strip_tatweel = st.sidebar.checkbox("Strip tatweel", True)
+        insert_white_spaces = st.sidebar.checkbox("Insert white spaces", True)
+        remove_non_digit_repetition = st.sidebar.checkbox(
+            "Remove non-digit repetition", True
+        )
+        replace_slash_with_dash = st.sidebar.checkbox("Replace slash with dash", None)
+        map_hindi_numbers_to_arabic = st.sidebar.checkbox(
+            "Map hindi numbers to arabic", None
+        )
+        apply_farasa_segmentation = st.sidebar.checkbox(
+            "Apply farasa segmentation", None
+        )
+
+    run_preprocessor = st.button("Run Pre-Processor")
+
+    prep_text = None
+    if run_preprocessor:
+        if model_selector == "None":
+            arabert_preprocessor = ArabertPreprocessor(
+                model_selector,
+                keep_emojis,
+                remove_html_markup,
+                replace_urls_emails_mentions,
+                strip_tashkeel,
+                strip_tatweel,
+                insert_white_spaces,
+                remove_non_digit_repetition,
+                replace_slash_with_dash,
+                map_hindi_numbers_to_arabic,
+                apply_farasa_segmentation,
+            )
+        else:
+            arabert_preprocessor = ArabertPreprocessor(model_name=model_selector)
+        prep_text = arabert_preprocessor._preprocess_v3(input_text)
+        st.write(prep_text)
+
+    st.write("-----")
+    input_text_unprep = st.text_input(
+        "Text to Undo the Pre-Processing",
+        value=prep_text
+        if prep_text
+        else "و+ لن نبالغ إذا قل +نا : إن ' هاتف ' أو ' كمبيوتر ال+ مكتب ' في زمن +نا هذا ضروري",
+    )
+    run_unpreprocessor = st.button("Run Un-Pre-Processor")
+
+    if run_unpreprocessor:
+        st.write(unpreprocess(input_text_unprep))

backend/qa.py

@@ -0,0 +1,50 @@
+import streamlit as st
+
+from .qa_utils import annotate_answer
+from .services import get_qa_answers
+
+
+def write():
+    _, col1, _ = st.columns(3)
+
+    with col1:
+        st.image("images/is2alni_logo.png", width=200)
+        st.title("إسألني أي شيء")
+
+    st.markdown(
+        """
+    <style>
+    p, div, input, label {
+    text-align: right;
+    }
+    </style>
+        """,
+        unsafe_allow_html=True,
+    )
+
+    st.sidebar.header("Info")
+    st.sidebar.image("images/AraELECTRA.png", width=150)
+    st.sidebar.write("Powered by [AraELECTRA](https://github.com/aub-mind/arabert)")
+
+    st.sidebar.write("\n")
+    n_answers = st.sidebar.slider(
+        "Max. number of answers", min_value=1, max_value=10, value=2, step=1
+    )
+
+    question = st.text_input("", value="من هو جو بايدن؟")
+    if "؟" not in question:
+        question += "؟"
+
+    run_query = st.button("أجب")
+    if run_query:
+        # https://discuss.streamlit.io/t/showing-a-gif-while-st-spinner-runs/5084
+        with st.spinner("... جاري البحث "):
+            results_dict = get_qa_answers(question)
+
+        if len(results_dict) > 0:
+            st.write("## :الأجابات هي")
+            for result in results_dict["results"][:n_answers]:
+                annotate_answer(result)
+                f"[**المصدر**](<{result['link']}>)"
+        else:
+            st.write("## 😞 ليس لدي جواب")

backend/qa_utils.py

@@ -0,0 +1,163 @@
+import streamlit.components.v1
+
+from htbuilder import HtmlElement, div, span, styles
+from htbuilder.units import px, rem, em
+
+
+def annotation(body, label="", background="#ddd", color="#333", **style):
+    """Build an HtmlElement span object with the given body and annotation label.
+
+    The end result will look something like this:
+
+        [body | label]
+
+    Parameters
+    ----------
+    body : string
+        The string to put in the "body" part of the annotation.
+    label : string
+        The string to put in the "label" part of the annotation.
+    background : string
+        The color to use for the background "chip" containing this annotation.
+    color : string
+        The color to use for the body and label text.
+    **style : dict
+        Any CSS you want to use to customize the containing "chip".
+
+    Examples
+    --------
+
+    Produce a simple annotation with default colors:
+
+    >>> annotation("apple", "fruit")
+
+    Produce an annotation with custom colors:
+
+    >>> annotation("apple", "fruit", background="#FF0", color="black")
+
+    Produce an annotation with crazy CSS:
+
+    >>> annotation("apple", "fruit", background="#FF0", border="1px dashed red")
+
+    """
+
+    if "font_family" not in style:
+        style["font_family"] = "sans-serif"
+
+    return span(
+        style=styles(
+            background=background,
+            border_radius=rem(0.33),
+            color=color,
+            padding=(rem(0.17), rem(0.67)),
+            display="inline-flex",
+            justify_content="center",
+            align_items="center",
+            **style,
+        )
+    )(
+        body,
+        span(
+            style=styles(
+                color=color,
+                font_size=em(0.67),
+                opacity=0.5,
+                padding_left=rem(0.5),
+                text_transform="uppercase",
+                margin_bottom=px(-2),
+            )
+        )(label),
+    )
+
+
+def annotated_text(*args, **kwargs):
+    """Writes test with annotations into your Streamlit app.
+
+    Parameters
+    ----------
+    *args : str, tuple or htbuilder.HtmlElement
+        Arguments can be:
+        - strings, to draw the string as-is on the screen.
+        - tuples of the form (main_text, annotation_text, background, color) where
+          background and foreground colors are optional and should be an CSS-valid string such as
+          "#aabbcc" or "rgb(10, 20, 30)"
+        - HtmlElement objects in case you want to customize the annotations further. In particular,
+          you can import the `annotation()` function from this module to easily produce annotations
+          whose CSS you can customize via keyword arguments.
+
+    Examples
+    --------
+
+    >>> annotated_text(
+    ...     "This ",
+    ...     ("is", "verb", "#8ef"),
+    ...     " some ",
+    ...     ("annotated", "adj", "#faa"),
+    ...     ("text", "noun", "#afa"),
+    ...     " for those of ",
+    ...     ("you", "pronoun", "#fea"),
+    ...     " who ",
+    ...     ("like", "verb", "#8ef"),
+    ...     " this sort of ",
+    ...     ("thing", "noun", "#afa"),
+    ... )
+
+    >>> annotated_text(
+    ...     "Hello ",
+    ...     annotation("world!", "noun", color="#8ef", border="1px dashed red"),
+    ... )
+
+    """
+    out = div(
+        style=styles(
+            font_family="sans-serif",
+            line_height="1.45",
+            font_size=px(16),
+            text_align="right",
+        )
+    )
+
+    for arg in args:
+        if isinstance(arg, str):
+            out(arg)
+
+        elif isinstance(arg, HtmlElement):
+            out(arg)
+
+        elif isinstance(arg, tuple):
+            out(annotation(*arg))
+
+        else:
+            raise Exception("Oh noes!")
+
+    streamlit.components.v1.html(str(out), **kwargs)
+
+
+def shorten_text(text, n, reverse=False):
+    if text.isspace() or text == "":
+        return text
+    if reverse:
+        text = text[::-1]
+    words = iter(text.split())
+    lines, current = [], next(words)
+    for word in words:
+        if len(current) + 1 + len(word) > n:
+            break
+        else:
+            current += " " + word
+    lines.append(current)
+    if reverse:
+        return lines[0][::-1]
+    return lines[0]
+
+
+def annotate_answer(result):
+    annotated_text(
+        shorten_text(
+            result["original"][: result["new_start"]],
+            500,
+            reverse=True,
+        ),
+        (result["new_answer"], "جواب", "#8ef"),
+        shorten_text(result["original"][result["new_end"] :], 500) + " ...... إلخ",
+    )

backend/sa.py

@@ -0,0 +1,76 @@
+import streamlit as st
+from .services import SentimentAnalyzer
+from functools import lru_cache
+
+# @st.cache(allow_output_mutation=False, hash_funcs={Tokenizer: str})
+@lru_cache(maxsize=1)
+def load_text_generator():
+    predictor = SentimentAnalyzer()
+    return predictor
+
+
+predictor = load_text_generator()
+
+
+def write():
+    st.markdown(
+        """
+        # Arabic Sentiment Analysis
+
+        This is a simple sentiment analysis app that uses the prediction kernel from Wissam's (me) submission that won the [Arabic Senitment Analysis competition @ KAUST](https://www.kaggle.com/c/arabic-sentiment-analysis-2021-kaust)
+        """
+    )
+    if st.checkbox("More info: "):
+        st.markdown(
+            """
+            ### Submission Description:
+
+            My submission is based on an ensemble of 5 models with varying preprocessing, and classifier design. All model variants are built over MARBERT [1], which is a BERT-based model pre-trained on 1B dialectal Arabic tweets.
+
+            For preprocessing, all models shared the following steps:
+            -	Replacing user mentions with “USER” and links with “URL”
+            -	Replacing the “#” with “HASH”
+            -	Removed the underscore character since it is missing the MARBERT vocabulary.
+            -	Removed diacritics and elongations (tatweel)
+            -	Spacing out emojis
+
+            For classifier design, all models use a dense layer on top of MARBERT unless otherwise specified. Model training is done by hyperparameter grid-search with 5-fold cross-validation with the following search space:
+            -	Learning rate: [2e-5,3e-5,4e-5]
+            -	Batch size: 128
+            -	Maximum sequence length: 64
+            -	Epochs: 3 (we select the best epoch for the final prediction)
+            -	Warmup ratio: [0,0.1]
+            -	Seed: [1,25,42,123,666]
+
+            Model I is a vanilla variant with only the preprocessing steps mention above applied. Model II enhances the emoji representation by replacing OOV emojis with ones that have similar meaning, for example 💊  😷.
+            We noticed the repetitive use of “السلام عليكم” and “ورحمة الله وبركاته” in neutral tweets, especially when users were directing questions to business accounts. This could confuse the classifier, if it encountered these words in a for example a negative tweet, hence in Model III we removed variation of the phrase mentioned before using fuzzy matching algorithms.
+
+            In Model IV, we tried to help the model by appending a sarcasm label to the input. We first trained a separate MARBERT on the ArSarcasm [2] dataset, and then used it to label the training and test sets.
+
+            Model V uses the vanilla preprocessing approach, but instead of a dense layer built on top of MARBERT, we follow the approach detailed by Safaya et.al. [3] which uses a CNN-based classifier instead.
+
+            For the final prediction, we first average the predictions of the 5 models from cross-validation (this is done for each model separately), we then average the results from the 5 model variants. We observed that the distribution of the predicted sentiment classes, doesn’t quite match the true distribution, this is due to the model preferring the neutral class over the positive class. To counter that, we apply what we call Label-Weighted average where during after the final averaging we rescale the score with the following weights 1.57,0.98 and 0.93 for positive, neutral, and negative (note that the weights were determined empirically).
+
+            1- https://aclanthology.org/2021.acl-long.551/
+
+            2- https://github.com/iabufarha/ArSarcasm
+
+            3- https://github.com/alisafaya/OffensEval2020
+
+
+            """
+        )
+    input_text = st.text_input(
+        "Enter your text here:",
+    )
+    if st.button("Predict"):
+        with st.spinner("Predicting..."):
+            prediction, score, all_score = predictor.predict([input_text])
+            st.write(f"Result: {prediction[0]}")
+            detailed_score = {
+                "Positive": all_score[0][0],
+                "Neutral": all_score[0][1],
+                "Negative": all_score[0][2],
+            }
+            st.write("All scores:")
+            st.write(detailed_score)

backend/sa_utils.py

@@ -0,0 +1,510 @@
+import re
+from contextlib import contextmanager
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from fuzzysearch import find_near_matches
+from pyarabic import araby
+from torch import nn
+from transformers import AutoTokenizer, BertModel, BertPreTrainedModel, pipeline
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+from .preprocess import ArabertPreprocessor, url_regexes, user_mention_regex
+
+multiple_char_pattern = re.compile(r"(.)\1{2,}", re.DOTALL)
+
+# ASAD-NEW_AraBERT_PREP-Balanced
+class NewArabicPreprocessorBalanced(ArabertPreprocessor):
+    def __init__(
+        self,
+        model_name: str,
+        keep_emojis: bool = False,
+        remove_html_markup: bool = True,
+        replace_urls_emails_mentions: bool = True,
+        strip_tashkeel: bool = True,
+        strip_tatweel: bool = True,
+        insert_white_spaces: bool = True,
+        remove_non_digit_repetition: bool = True,
+        replace_slash_with_dash: bool = None,
+        map_hindi_numbers_to_arabic: bool = None,
+        apply_farasa_segmentation: bool = None,
+    ):
+        if "UBC-NLP" in model_name or "CAMeL-Lab" in model_name:
+            keep_emojis = True
+            remove_non_digit_repetition = True
+        super().__init__(
+            model_name=model_name,
+            keep_emojis=keep_emojis,
+            remove_html_markup=remove_html_markup,
+            replace_urls_emails_mentions=replace_urls_emails_mentions,
+            strip_tashkeel=strip_tashkeel,
+            strip_tatweel=strip_tatweel,
+            insert_white_spaces=insert_white_spaces,
+            remove_non_digit_repetition=remove_non_digit_repetition,
+            replace_slash_with_dash=replace_slash_with_dash,
+            map_hindi_numbers_to_arabic=map_hindi_numbers_to_arabic,
+            apply_farasa_segmentation=apply_farasa_segmentation,
+        )
+        self.true_model_name = model_name
+
+    def preprocess(self, text):
+        if "UBC-NLP" in self.true_model_name:
+            return self.ubc_prep(text)
+
+    def ubc_prep(self, text):
+        text = re.sub("\s", " ", text)
+        text = text.replace("\\n", " ")
+        text = text.replace("\\r", " ")
+        text = araby.strip_tashkeel(text)
+        text = araby.strip_tatweel(text)
+        # replace all possible URLs
+        for reg in url_regexes:
+            text = re.sub(reg, " URL ", text)
+        text = re.sub("(URL\s*)+", " URL ", text)
+        # replace mentions with USER
+        text = re.sub(user_mention_regex, " USER ", text)
+        text = re.sub("(USER\s*)+", " USER ", text)
+        # replace hashtags with HASHTAG
+        # text = re.sub(r"#[\w\d]+", " HASH TAG ", text)
+        text = text.replace("#", " HASH ")
+        text = text.replace("_", " ")
+        text = " ".join(text.split())
+        # text = re.sub("\B\\[Uu]\w+", "", text)
+        text = text.replace("\\U0001f97a", "🥺")
+        text = text.replace("\\U0001f928", "🤨")
+        text = text.replace("\\U0001f9d8", "😀")
+        text = text.replace("\\U0001f975", "😥")
+        text = text.replace("\\U0001f92f", "😲")
+        text = text.replace("\\U0001f92d", "🤭")
+        text = text.replace("\\U0001f9d1", "😐")
+        text = text.replace("\\U000e0067", "")
+        text = text.replace("\\U000e006e", "")
+        text = text.replace("\\U0001f90d", "♥")
+        text = text.replace("\\U0001f973", "🎉")
+        text = text.replace("\\U0001fa79", "")
+        text = text.replace("\\U0001f92b", "🤐")
+        text = text.replace("\\U0001f9da", "🦋")
+        text = text.replace("\\U0001f90e", "♥")
+        text = text.replace("\\U0001f9d0", "🧐")
+        text = text.replace("\\U0001f9cf", "")
+        text = text.replace("\\U0001f92c", "😠")
+        text = text.replace("\\U0001f9f8", "😸")
+        text = text.replace("\\U0001f9b6", "💩")
+        text = text.replace("\\U0001f932", "🤲")
+        text = text.replace("\\U0001f9e1", "🧡")
+        text = text.replace("\\U0001f974", "☹")
+        text = text.replace("\\U0001f91f", "")
+        text = text.replace("\\U0001f9fb", "💩")
+        text = text.replace("\\U0001f92a", "🤪")
+        text = text.replace("\\U0001f9fc", "")
+        text = text.replace("\\U000e0065", "")
+        text = text.replace("\\U0001f92e", "💩")
+        text = text.replace("\\U000e007f", "")
+        text = text.replace("\\U0001f970", "🥰")
+        text = text.replace("\\U0001f929", "🤩")
+        text = text.replace("\\U0001f6f9", "")
+        text = text.replace("🤍", "♥")
+        text = text.replace("🦠", "😷")
+        text = text.replace("🤢", "مقرف")
+        text = text.replace("🤮", "مقرف")
+        text = text.replace("🕠", "⌚")
+        text = text.replace("🤬", "😠")
+        text = text.replace("🤧", "😷")
+        text = text.replace("🥳", "🎉")
+        text = text.replace("🥵", "🔥")
+        text = text.replace("🥴", "☹")
+        text = text.replace("🤫", "🤐")
+        text = text.replace("🤥", "كذاب")
+        text = text.replace("\\u200d", " ")
+        text = text.replace("u200d", " ")
+        text = text.replace("\\u200c", " ")
+        text = text.replace("u200c", " ")
+        text = text.replace('"', "'")
+        text = text.replace("\\xa0", "")
+        text = text.replace("\\u2066", " ")
+        text = re.sub("\B\\\[Uu]\w+", "", text)
+        text = super(NewArabicPreprocessorBalanced, self).preprocess(text)
+
+        text = " ".join(text.split())
+        return text
+
+
+"""CNNMarbertArabicPreprocessor"""
+# ASAD-CNN_MARBERT
+class CNNMarbertArabicPreprocessor(ArabertPreprocessor):
+    def __init__(
+        self,
+        model_name,
+        keep_emojis=False,
+        remove_html_markup=True,
+        replace_urls_emails_mentions=True,
+        remove_elongations=True,
+    ):
+        if "UBC-NLP" in model_name or "CAMeL-Lab" in model_name:
+            keep_emojis = True
+            remove_elongations = False
+        super().__init__(
+            model_name,
+            keep_emojis,
+            remove_html_markup,
+            replace_urls_emails_mentions,
+            remove_elongations,
+        )
+        self.true_model_name = model_name
+
+    def preprocess(self, text):
+        if "UBC-NLP" in self.true_model_name:
+            return self.ubc_prep(text)
+
+    def ubc_prep(self, text):
+        text = re.sub("\s", " ", text)
+        text = text.replace("\\n", " ")
+        text = araby.strip_tashkeel(text)
+        text = araby.strip_tatweel(text)
+        # replace all possible URLs
+        for reg in url_regexes:
+            text = re.sub(reg, " URL ", text)
+        text = re.sub("(URL\s*)+", " URL ", text)
+        # replace mentions with USER
+        text = re.sub(user_mention_regex, " USER ", text)
+        text = re.sub("(USER\s*)+", " USER ", text)
+        # replace hashtags with HASHTAG
+        # text = re.sub(r"#[\w\d]+", " HASH TAG ", text)
+        text = text.replace("#", " HASH ")
+        text = text.replace("_", " ")
+        text = " ".join(text.split())
+        text = super(CNNMarbertArabicPreprocessor, self).preprocess(text)
+        text = text.replace("\u200d", " ")
+        text = text.replace("u200d", " ")
+        text = text.replace("\u200c", " ")
+        text = text.replace("u200c", " ")
+        text = text.replace('"', "'")
+        # text = re.sub('[\d\.]+', ' NUM ', text)
+        # text = re.sub('(NUM\s*)+', ' NUM ', text)
+        text = multiple_char_pattern.sub(r"\1\1", text)
+        text = " ".join(text.split())
+        return text
+
+
+"""Trial5ArabicPreprocessor"""
+
+
+class Trial5ArabicPreprocessor(ArabertPreprocessor):
+    def __init__(
+        self,
+        model_name,
+        keep_emojis=False,
+        remove_html_markup=True,
+        replace_urls_emails_mentions=True,
+    ):
+        if "UBC-NLP" in model_name:
+            keep_emojis = True
+        super().__init__(
+            model_name, keep_emojis, remove_html_markup, replace_urls_emails_mentions
+        )
+        self.true_model_name = model_name
+
+    def preprocess(self, text):
+        if "UBC-NLP" in self.true_model_name:
+            return self.ubc_prep(text)
+
+    def ubc_prep(self, text):
+        text = re.sub("\s", " ", text)
+        text = text.replace("\\n", " ")
+        text = araby.strip_tashkeel(text)
+        text = araby.strip_tatweel(text)
+        # replace all possible URLs
+        for reg in url_regexes:
+            text = re.sub(reg, " URL ", text)
+        # replace mentions with USER
+        text = re.sub(user_mention_regex, " USER ", text)
+        # replace hashtags with HASHTAG
+        # text = re.sub(r"#[\w\d]+", " HASH TAG ", text)
+        text = text.replace("#", " HASH TAG ")
+        text = text.replace("_", " ")
+        text = " ".join(text.split())
+        text = super(Trial5ArabicPreprocessor, self).preprocess(text)
+        # text = text.replace("السلام عليكم"," ")
+        # text = text.replace(find_near_matches("السلام عليكم",text,max_deletions=3,max_l_dist=3)[0].matched," ")
+        return text
+
+
+"""SarcasmArabicPreprocessor"""
+
+
+class SarcasmArabicPreprocessor(ArabertPreprocessor):
+    def __init__(
+        self,
+        model_name,
+        keep_emojis=False,
+        remove_html_markup=True,
+        replace_urls_emails_mentions=True,
+    ):
+        if "UBC-NLP" in model_name:
+            keep_emojis = True
+        super().__init__(
+            model_name, keep_emojis, remove_html_markup, replace_urls_emails_mentions
+        )
+        self.true_model_name = model_name
+
+    def preprocess(self, text):
+        if "UBC-NLP" in self.true_model_name:
+            return self.ubc_prep(text)
+        else:
+            return super(SarcasmArabicPreprocessor, self).preprocess(text)
+
+    def ubc_prep(self, text):
+        text = re.sub("\s", " ", text)
+        text = text.replace("\\n", " ")
+        text = araby.strip_tashkeel(text)
+        text = araby.strip_tatweel(text)
+        # replace all possible URLs
+        for reg in url_regexes:
+            text = re.sub(reg, " URL ", text)
+        # replace mentions with USER
+        text = re.sub(user_mention_regex, " USER ", text)
+        # replace hashtags with HASHTAG
+        # text = re.sub(r"#[\w\d]+", " HASH TAG ", text)
+        text = text.replace("#", " HASH TAG ")
+        text = text.replace("_", " ")
+        text = text.replace('"', " ")
+        text = " ".join(text.split())
+        text = super(SarcasmArabicPreprocessor, self).preprocess(text)
+        return text
+
+
+"""NoAOAArabicPreprocessor"""
+
+
+class NoAOAArabicPreprocessor(ArabertPreprocessor):
+    def __init__(
+        self,
+        model_name,
+        keep_emojis=False,
+        remove_html_markup=True,
+        replace_urls_emails_mentions=True,
+    ):
+        if "UBC-NLP" in model_name:
+            keep_emojis = True
+        super().__init__(
+            model_name, keep_emojis, remove_html_markup, replace_urls_emails_mentions
+        )
+        self.true_model_name = model_name
+
+    def preprocess(self, text):
+        if "UBC-NLP" in self.true_model_name:
+            return self.ubc_prep(text)
+        else:
+            return super(NoAOAArabicPreprocessor, self).preprocess(text)
+
+    def ubc_prep(self, text):
+        text = re.sub("\s", " ", text)
+        text = text.replace("\\n", " ")
+        text = araby.strip_tashkeel(text)
+        text = araby.strip_tatweel(text)
+        # replace all possible URLs
+        for reg in url_regexes:
+            text = re.sub(reg, " URL ", text)
+        # replace mentions with USER
+        text = re.sub(user_mention_regex, " USER ", text)
+        # replace hashtags with HASHTAG
+        # text = re.sub(r"#[\w\d]+", " HASH TAG ", text)
+        text = text.replace("#", " HASH TAG ")
+        text = text.replace("_", " ")
+        text = " ".join(text.split())
+        text = super(NoAOAArabicPreprocessor, self).preprocess(text)
+        text = text.replace("السلام عليكم", " ")
+        text = text.replace("ورحمة الله وبركاته", " ")
+        matched = find_near_matches("السلام عليكم", text, max_deletions=3, max_l_dist=3)
+        if len(matched) > 0:
+            text = text.replace(matched[0].matched, " ")
+        matched = find_near_matches(
+            "ورحمة الله وبركاته", text, max_deletions=3, max_l_dist=3
+        )
+        if len(matched) > 0:
+            text = text.replace(matched[0].matched, " ")
+        return text
+
+
+class CnnBertForSequenceClassification(BertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.config = config
+
+        self.bert = BertModel(config)
+
+        filter_sizes = [1, 2, 3, 4, 5]
+        num_filters = 32
+        self.convs1 = nn.ModuleList(
+            [nn.Conv2d(4, num_filters, (K, config.hidden_size)) for K in filter_sizes]
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(len(filter_sizes) * num_filters, config.num_labels)
+
+        self.init_weights()
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        x = outputs[2][-4:]
+
+        x = torch.stack(x, dim=1)
+        x = [F.relu(conv(x)).squeeze(3) for conv in self.convs1]
+        x = [F.max_pool1d(i, i.size(2)).squeeze(2) for i in x]
+        x = torch.cat(x, 1)
+        x = self.dropout(x)
+        logits = self.classifier(x)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (
+                    labels.dtype == torch.long or labels.dtype == torch.int
+                ):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = nn.MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = nn.CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = nn.BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=None,
+            attentions=outputs.attentions,
+        )
+
+
+class CNNTextClassificationPipeline:
+    def __init__(self, model_path, device, return_all_scores=False):
+        self.model_path = model_path
+        self.model = CnnBertForSequenceClassification.from_pretrained(self.model_path)
+        # Special handling
+        self.device = torch.device("cpu" if device < 0 else f"cuda:{device}")
+        if self.device.type == "cuda":
+            self.model = self.model.to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.return_all_scores = return_all_scores
+
+    @contextmanager
+    def device_placement(self):
+        """
+        Context Manager allowing tensor allocation on the user-specified device in framework agnostic way.
+        Returns:
+            Context manager
+        Examples::
+            # Explicitly ask for tensor allocation on CUDA device :0
+            pipe = pipeline(..., device=0)
+            with pipe.device_placement():
+                # Every framework specific tensor allocation will be done on the request device
+                output = pipe(...)
+        """
+
+        if self.device.type == "cuda":
+            torch.cuda.set_device(self.device)
+
+        yield
+
+    def ensure_tensor_on_device(self, **inputs):
+        """
+        Ensure PyTorch tensors are on the specified device.
+        Args:
+            inputs (keyword arguments that should be :obj:`torch.Tensor`): The tensors to place on :obj:`self.device`.
+        Return:
+            :obj:`Dict[str, torch.Tensor]`: The same as :obj:`inputs` but on the proper device.
+        """
+        return {
+            name: tensor.to(self.device) if isinstance(tensor, torch.Tensor) else tensor
+            for name, tensor in inputs.items()
+        }
+
+    def __call__(self, text):
+        """
+        Classify the text(s) given as inputs.
+        Args:
+            args (:obj:`str` or :obj:`List[str]`):
+                One or several texts (or one list of prompts) to classify.
+        Return:
+            A list or a list of list of :obj:`dict`: Each result comes as list of dictionaries with the following keys:
+            - **label** (:obj:`str`) -- The label predicted.
+            - **score** (:obj:`float`) -- The corresponding probability.
+            If ``self.return_all_scores=True``, one such dictionary is returned per label.
+        """
+        # outputs = super().__call__(*args, **kwargs)
+        inputs = self.tokenizer.batch_encode_plus(
+            text,
+            add_special_tokens=True,
+            max_length=64,
+            padding=True,
+            truncation="longest_first",
+            return_tensors="pt",
+        )
+
+        with torch.no_grad():
+            inputs = self.ensure_tensor_on_device(**inputs)
+            predictions = self.model(**inputs)[0].cpu()
+
+        predictions = predictions.numpy()
+
+        if self.model.config.num_labels == 1:
+            scores = 1.0 / (1.0 + np.exp(-predictions))
+        else:
+            scores = np.exp(predictions) / np.exp(predictions).sum(-1, keepdims=True)
+        if self.return_all_scores:
+            return [
+                [
+                    {"label": self.model.config.id2label[i], "score": score.item()}
+                    for i, score in enumerate(item)
+                ]
+                for item in scores
+            ]
+        else:
+            return [
+                {"label": self.inv_label_map[item.argmax()], "score": item.max().item()}
+                for item in scores
+            ]

backend/sarcasm.py

@@ -0,0 +1,21 @@
+import streamlit as st
+from .sa import predictor
+
+
+def write():
+    st.markdown(
+        """
+        # Arabic Sarcasm Detection
+
+        This is a simple sarcasm detection app that uses the [MARBERT](https://huggingface.co/UBC-NLP/MARBERT) model trained on [ArSarcasm](https://github.com/iabufarha/ArSarcasm)
+        """
+    )
+
+    input_text = st.text_input(
+        "Enter your text here:",
+    )
+    if st.button("Predict"):
+        with st.spinner("Predicting..."):
+            prediction, scores = predictor.get_preds_from_sarcasm([input_text])
+            st.write(f"Result: {prediction[0]}")
+            st.write(f"Score: {scores[0]}")

backend/services.py

@@ -0,0 +1,519 @@
+import json
+import logging
+import os
+from functools import lru_cache
+from typing import List
+from urllib.parse import unquote
+
+import more_itertools
+import pandas as pd
+import requests
+import streamlit as st
+import wikipedia
+from codetiming import Timer
+from fuzzysearch import find_near_matches
+from googleapi import google
+from tqdm.auto import tqdm
+from transformers import (
+    AutoTokenizer,
+    GPT2LMHeadModel,
+    GPT2Tokenizer,
+    pipeline,
+    set_seed,
+)
+
+from .modeling_gpt2 import GPT2LMHeadModel as GROVERLMHeadModel
+from .preprocess import ArabertPreprocessor
+from .sa_utils import *
+from .utils import download_models, softmax
+
+logger = logging.getLogger(__name__)
+# Taken and Modified from https://huggingface.co/spaces/flax-community/chef-transformer/blob/main/app.py
+class TextGeneration:
+    def __init__(self):
+        self.debug = False
+        self.generation_pipline = {}
+        self.preprocessor = ArabertPreprocessor(model_name="aragpt2-mega")
+        self.tokenizer = GPT2Tokenizer.from_pretrained(
+            "aubmindlab/aragpt2-mega", use_fast=False
+        )
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.API_KEY = os.getenv("API_KEY")
+        self.headers = {"Authorization": f"Bearer {self.API_KEY}"}
+        # self.model_names_or_paths = {
+        #     "aragpt2-medium": "D:/ML/Models/aragpt2-medium",
+        #     "aragpt2-base": "D:/ML/Models/aragpt2-base",
+        # }
+        self.model_names_or_paths = {
+            # "aragpt2-medium": "aubmindlab/aragpt2-medium",
+            "aragpt2-base": "aubmindlab/aragpt2-base",
+            # "aragpt2-large": "aubmindlab/aragpt2-large",
+            "aragpt2-mega": "aubmindlab/aragpt2-mega",
+        }
+        set_seed(42)
+
+    def load_pipeline(self):
+        for model_name, model_path in self.model_names_or_paths.items():
+            if "base" in model_name or "medium" in model_name:
+                self.generation_pipline[model_name] = pipeline(
+                    "text-generation",
+                    model=GPT2LMHeadModel.from_pretrained(model_path),
+                    tokenizer=self.tokenizer,
+                    device=-1,
+                )
+            else:
+                self.generation_pipline[model_name] = pipeline(
+                    "text-generation",
+                    model=GROVERLMHeadModel.from_pretrained(model_path),
+                    tokenizer=self.tokenizer,
+                    device=-1,
+                )
+
+    def load(self):
+        if not self.debug:
+            self.load_pipeline()
+
+    def generate(
+        self,
+        model_name,
+        prompt,
+        max_new_tokens: int,
+        temperature: float,
+        top_k: int,
+        top_p: float,
+        repetition_penalty: float,
+        no_repeat_ngram_size: int,
+        do_sample: bool,
+        num_beams: int,
+    ):
+        logger.info(f"Generating with {model_name}")
+        prompt = self.preprocessor.preprocess(prompt)
+        return_full_text = False
+        return_text = True
+        num_return_sequences = 1
+        pad_token_id = 0
+        eos_token_id = 0
+        input_tok = self.tokenizer.tokenize(prompt)
+        max_length = len(input_tok) + max_new_tokens
+        if max_length > 1024:
+            max_length = 1024
+        if not self.debug:
+            generated_text = self.generation_pipline[model_name.lower()](
+                prompt,
+                max_length=max_length,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                no_repeat_ngram_size=no_repeat_ngram_size,
+                pad_token_id=pad_token_id,
+                eos_token_id=eos_token_id,
+                return_full_text=return_full_text,
+                return_text=return_text,
+                do_sample=do_sample,
+                num_beams=num_beams,
+                num_return_sequences=num_return_sequences,
+            )[0]["generated_text"]
+        else:
+            generated_text = self.generate_by_query(
+                prompt,
+                model_name,
+                max_length=max_length,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                no_repeat_ngram_size=no_repeat_ngram_size,
+                pad_token_id=pad_token_id,
+                eos_token_id=eos_token_id,
+                return_full_text=return_full_text,
+                return_text=return_text,
+                do_sample=do_sample,
+                num_beams=num_beams,
+                num_return_sequences=num_return_sequences,
+            )
+            # print(generated_text)
+            if isinstance(generated_text, dict):
+                if "error" in generated_text:
+                    if "is currently loading" in generated_text["error"]:
+                        return f"Model is currently loading, estimated time is {generated_text['estimated_time']}"
+                    return generated_text["error"]
+                else:
+                    return "Something happened 🤷‍♂️!!"
+            else:
+                generated_text = generated_text[0]["generated_text"]
+
+        logger.info(f"Prompt: {prompt}")
+        logger.info(f"Generated text: {generated_text}")
+        return self.preprocessor.unpreprocess(generated_text)
+
+    def query(self, payload, model_name):
+        data = json.dumps(payload)
+        url = (
+            "https://api-inference.huggingface.co/models/aubmindlab/"
+            + model_name.lower()
+        )
+        response = requests.request("POST", url, headers=self.headers, data=data)
+        return json.loads(response.content.decode("utf-8"))
+
+    def generate_by_query(
+        self,
+        prompt: str,
+        model_name: str,
+        max_length: int,
+        temperature: float,
+        top_k: int,
+        top_p: float,
+        repetition_penalty: float,
+        no_repeat_ngram_size: int,
+        pad_token_id: int,
+        eos_token_id: int,
+        return_full_text: int,
+        return_text: int,
+        do_sample: bool,
+        num_beams: int,
+        num_return_sequences: int,
+    ):
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "max_length ": max_length,
+                "top_k": top_k,
+                "top_p": top_p,
+                "temperature": temperature,
+                "repetition_penalty": repetition_penalty,
+                "no_repeat_ngram_size": no_repeat_ngram_size,
+                "pad_token_id": pad_token_id,
+                "eos_token_id": eos_token_id,
+                "return_full_text": return_full_text,
+                "return_text": return_text,
+                "pad_token_id": pad_token_id,
+                "do_sample": do_sample,
+                "num_beams": num_beams,
+                "num_return_sequences": num_return_sequences,
+            },
+            "options": {
+                "use_cache": True,
+            },
+        }
+        return self.query(payload, model_name)
+
+
+class SentimentAnalyzer:
+    def __init__(self):
+        self.sa_models = [
+            "sa_trial5_1",
+            # "sa_no_aoa_in_neutral",
+            # "sa_cnnbert",
+            # "sa_sarcasm",
+            # "sar_trial10",
+            # "sa_no_AOA",
+        ]
+        download_models(self.sa_models)
+        # fmt: off
+        self.processors = {
+            "sa_trial5_1": Trial5ArabicPreprocessor(model_name='UBC-NLP/MARBERT'),
+            # "sa_no_aoa_in_neutral": NewArabicPreprocessorBalanced(model_name='UBC-NLP/MARBERT'),
+            # "sa_cnnbert": CNNMarbertArabicPreprocessor(model_name='UBC-NLP/MARBERT'),
+            # "sa_sarcasm": SarcasmArabicPreprocessor(model_name='UBC-NLP/MARBERT'),
+            # "sar_trial10": SarcasmArabicPreprocessor(model_name='UBC-NLP/MARBERT'),
+            # "sa_no_AOA": NewArabicPreprocessorBalanced(model_name='UBC-NLP/MARBERT'),
+        }
+
+        self.pipelines = {
+            "sa_trial5_1": [pipeline("sentiment-analysis", model="{}/train_{}/best_model".format("sa_trial5_1",i), device=-1,return_all_scores =True) for i in tqdm(range(0,5), desc=f"Loading pipeline for model: sa_trial5_1")],
+            # "sa_no_aoa_in_neutral": [pipeline("sentiment-analysis", model="{}/train_{}/best_model".format("sa_no_aoa_in_neutral",i), device=-1,return_all_scores =True) for i in tqdm(range(0,5), desc=f"Loading pipeline for model: sa_no_aoa_in_neutral")],
+            # "sa_cnnbert": [CNNTextClassificationPipeline("{}/train_{}/best_model".format("sa_cnnbert",i), device=-1, return_all_scores =True) for i in tqdm(range(0,5), desc=f"Loading pipeline for model: sa_cnnbert")],
+            # "sa_sarcasm": [pipeline("sentiment-analysis", model="{}/train_{}/best_model".format("sa_sarcasm",i), device=-1,return_all_scores =True) for i in tqdm(range(0,5), desc=f"Loading pipeline for model: sa_sarcasm")],
+            # "sar_trial10": [pipeline("sentiment-analysis", model="{}/train_{}/best_model".format("sar_trial10",i), device=-1,return_all_scores =True) for i in tqdm(range(0,5), desc=f"Loading pipeline for model: sar_trial10")],
+            # "sa_no_AOA": [pipeline("sentiment-analysis", model="{}/train_{}/best_model".format("sa_no_AOA",i), device=-1,return_all_scores =True) for i in tqdm(range(0,5), desc=f"Loading pipeline for model: sa_no_AOA")],
+        }
+        # fmt: on
+
+    def get_preds_from_sarcasm(self, texts):
+        prep = self.processors["sar_trial10"]
+        prep_texts = [prep.preprocess(x) for x in texts]
+
+        preds_df = pd.DataFrame([])
+        for i in range(0, 5):
+            preds = []
+            for s in more_itertools.chunked(list(prep_texts), 128):
+                preds.extend(self.pipelines["sar_trial10"][i](s))
+            preds_df[f"model_{i}"] = preds
+
+        final_labels = []
+        final_scores = []
+        for id, row in preds_df.iterrows():
+            pos_total = 0
+            neu_total = 0
+            for pred in row[:]:
+                pos_total += pred[0]["score"]
+                neu_total += pred[1]["score"]
+
+            pos_avg = pos_total / len(row[:])
+            neu_avg = neu_total / len(row[:])
+
+            final_labels.append(
+                self.pipelines["sar_trial10"][0].model.config.id2label[
+                    np.argmax([pos_avg, neu_avg])
+                ]
+            )
+            final_scores.append(np.max([pos_avg, neu_avg]))
+
+        return final_labels, final_scores
+
+    def get_preds_from_a_model(self, texts: List[str], model_name):
+        try:
+            prep = self.processors[model_name]
+
+            prep_texts = [prep.preprocess(x) for x in texts]
+            if model_name == "sa_sarcasm":
+                sarcasm_label, _ = self.get_preds_from_sarcasm(texts)
+                sarcastic_map = {"Not_Sarcastic": "غير ساخر", "Sarcastic": "ساخر"}
+                labeled_prep_texts = []
+                for t, l in zip(prep_texts, sarcasm_label):
+                    labeled_prep_texts.append(sarcastic_map[l] + " [SEP] " + t)
+
+            preds_df = pd.DataFrame([])
+            for i in range(0, 5):
+                preds = []
+                for s in more_itertools.chunked(list(prep_texts), 128):
+                    preds.extend(self.pipelines[model_name][i](s))
+                preds_df[f"model_{i}"] = preds
+
+            final_labels = []
+            final_scores = []
+            final_scores_list = []
+            for id, row in preds_df.iterrows():
+                pos_total = 0
+                neg_total = 0
+                neu_total = 0
+                for pred in row[2:]:
+                    pos_total += pred[0]["score"]
+                    neu_total += pred[1]["score"]
+                    neg_total += pred[2]["score"]
+
+                pos_avg = pos_total / 5
+                neu_avg = neu_total / 5
+                neg_avg = neg_total / 5
+
+                if model_name == "sa_no_aoa_in_neutral":
+                    final_labels.append(
+                        self.pipelines[model_name][0].model.config.id2label[
+                            np.argmax([neu_avg, neg_avg, pos_avg])
+                        ]
+                    )
+                else:
+                    final_labels.append(
+                        self.pipelines[model_name][0].model.config.id2label[
+                            np.argmax([pos_avg, neu_avg, neg_avg])
+                        ]
+                    )
+                final_scores.append(np.max([pos_avg, neu_avg, neg_avg]))
+                final_scores_list.append((pos_avg, neu_avg, neg_avg))
+        except RuntimeError as e:
+            if model_name == "sa_cnnbert":
+                return (
+                    ["Neutral"] * len(texts),
+                    [0.0] * len(texts),
+                    [(0.0, 0.0, 0.0)] * len(texts),
+                )
+            else:
+                raise RuntimeError(e)
+        return final_labels, final_scores, final_scores_list
+
+    def predict(self, texts: List[str]):
+        logger.info(f"Predicting for: {texts}")
+        # (
+        #     new_balanced_label,
+        #     new_balanced_score,
+        #     new_balanced_score_list,
+        # ) = self.get_preds_from_a_model(texts, "sa_no_aoa_in_neutral")
+        # (
+        #     cnn_marbert_label,
+        #     cnn_marbert_score,
+        #     cnn_marbert_score_list,
+        # ) = self.get_preds_from_a_model(texts, "sa_cnnbert")
+        trial5_label, trial5_score, trial5_score_list = self.get_preds_from_a_model(
+            texts, "sa_trial5_1"
+        )
+        # no_aoa_label, no_aoa_score, no_aoa_score_list = self.get_preds_from_a_model(
+        #     texts, "sa_no_AOA"
+        # )
+        # sarcasm_label, sarcasm_score, sarcasm_score_list = self.get_preds_from_a_model(
+        #     texts, "sa_sarcasm"
+        # )
+
+        id_label_map = {0: "Positive", 1: "Neutral", 2: "Negative"}
+
+        final_ensemble_prediction = []
+        final_ensemble_score = []
+        final_ensemble_all_score = []
+        for entry in zip(
+            # new_balanced_score_list,
+            # cnn_marbert_score_list,
+            trial5_score_list,
+            # no_aoa_score_list,
+            # sarcasm_score_list,
+        ):
+            pos_score = 0
+            neu_score = 0
+            neg_score = 0
+            for s in entry:
+                pos_score += s[0] * 1.57
+                neu_score += s[1] * 0.98
+                neg_score += s[2] * 0.93
+
+                # weighted 2
+                # pos_score += s[0]*1.67
+                # neu_score += s[1]
+                # neg_score += s[2]*0.95
+
+            final_ensemble_prediction.append(
+                id_label_map[np.argmax([pos_score, neu_score, neg_score])]
+            )
+            final_ensemble_score.append(np.max([pos_score, neu_score, neg_score]))
+            final_ensemble_all_score.append(
+                softmax(np.array([pos_score, neu_score, neg_score])).tolist()
+            )
+
+        logger.info(f"Result: {final_ensemble_prediction}")
+        logger.info(f"Score: {final_ensemble_score}")
+        logger.info(f"All Scores: {final_ensemble_all_score}")
+        return final_ensemble_prediction, final_ensemble_score, final_ensemble_all_score
+
+
+wikipedia.set_lang("ar")
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+preprocessor = ArabertPreprocessor("wissamantoun/araelectra-base-artydiqa")
+logger.info("Loading QA Pipeline...")
+tokenizer = AutoTokenizer.from_pretrained("wissamantoun/araelectra-base-artydiqa")
+qa_pipe = pipeline("question-answering", model="wissamantoun/araelectra-base-artydiqa")
+logger.info("Finished loading QA Pipeline...")
+
+
+@lru_cache(maxsize=100)
+def get_qa_answers(question):
+    logger.info("\n=================================================================")
+    logger.info(f"Question: {question}")
+
+    if "وسام أنطون" in question or "wissam antoun" in question.lower():
+        return {
+            "title": "Creator",
+            "results": [
+                {
+                    "score": 1.0,
+                    "new_start": 0,
+                    "new_end": 12,
+                    "new_answer": "My Creator 😜",
+                    "original": "My Creator 😜",
+                    "link": "https://github.com/WissamAntoun/",
+                }
+            ],
+        }
+    search_timer = Timer(
+        "search and wiki", text="Search and Wikipedia Time: {:.2f}", logger=logging.info
+    )
+    try:
+        search_timer.start()
+        search_results = google.search(
+            question + " site:ar.wikipedia.org", lang="ar", area="ar"
+        )
+        if len(search_results) == 0:
+            return {}
+
+        page_name = search_results[0].link.split("wiki/")[-1]
+        wiki_page = wikipedia.page(unquote(page_name))
+        wiki_page_content = wiki_page.content
+        search_timer.stop()
+    except:
+        return {}
+
+    sections = []
+    for section in re.split("== .+ ==[^=]", wiki_page_content):
+        if not section.isspace():
+            prep_section = tokenizer.tokenize(preprocessor.preprocess(section))
+            if len(prep_section) > 500:
+                subsections = []
+                for subsection in re.split("=== .+ ===", section):
+                    if subsection.isspace():
+                        continue
+                    prep_subsection = tokenizer.tokenize(
+                        preprocessor.preprocess(subsection)
+                    )
+                    subsections.append(subsection)
+                    # logger.info(f"Subsection found with length: {len(prep_subsection)}")
+                sections.extend(subsections)
+            else:
+                # logger.info(f"Regular Section with length: {len(prep_section)}")
+                sections.append(section)
+
+    full_len_sections = []
+    temp_section = ""
+    for section in sections:
+        if (
+            len(tokenizer.tokenize(preprocessor.preprocess(temp_section)))
+            + len(tokenizer.tokenize(preprocessor.preprocess(section)))
+            > 384
+        ):
+            if temp_section == "":
+                temp_section = section
+                continue
+            full_len_sections.append(temp_section)
+            # logger.info(
+            #     f"full section length: {len(tokenizer.tokenize(preprocessor.preprocess(temp_section)))}"
+            # )
+            temp_section = ""
+        else:
+            temp_section += " " + section + " "
+    if temp_section != "":
+        full_len_sections.append(temp_section)
+
+    reader_time = Timer("electra", text="Reader Time: {:.2f}", logger=logging.info)
+    reader_time.start()
+    results = qa_pipe(
+        question=[preprocessor.preprocess(question)] * len(full_len_sections),
+        context=[preprocessor.preprocess(x) for x in full_len_sections],
+    )
+
+    if not isinstance(results, list):
+        results = [results]
+
+    logger.info(f"Wiki Title: {unquote(page_name)}")
+    logger.info(f"Total Sections: {len(sections)}")
+    logger.info(f"Total Full Sections: {len(full_len_sections)}")
+
+    for result, section in zip(results, full_len_sections):
+        result["original"] = section
+        answer_match = find_near_matches(
+            " " + preprocessor.unpreprocess(result["answer"]) + " ",
+            result["original"],
+            max_l_dist=min(5, len(preprocessor.unpreprocess(result["answer"])) // 2),
+            max_deletions=0,
+        )
+        try:
+            result["new_start"] = answer_match[0].start
+            result["new_end"] = answer_match[0].end
+            result["new_answer"] = answer_match[0].matched
+            result["link"] = (
+                search_results[0].link + "#:~:text=" + result["new_answer"].strip()
+            )
+        except:
+            result["new_start"] = result["start"]
+            result["new_end"] = result["end"]
+            result["new_answer"] = result["answer"]
+            result["original"] = preprocessor.preprocess(result["original"])
+            result["link"] = search_results[0].link
+        logger.info(f"Answers: {preprocessor.preprocess(result['new_answer'])}")
+
+    sorted_results = sorted(results, reverse=True, key=lambda x: x["score"])
+
+    return_dict = {}
+    return_dict["title"] = unquote(page_name)
+    return_dict["results"] = sorted_results
+
+    reader_time.stop()
+    logger.info(f"Total time spent: {reader_time.last + search_timer.last}")
+    return return_dict

backend/utils.py

@@ -0,0 +1,64 @@
+import re
+import numpy as np
+import psutil
+import os
+from tqdm.auto import tqdm
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def get_current_ram_usage():
+    ram = psutil.virtual_memory()
+    return ram.available / 1024 / 1024 / 1024, ram.total / 1024 / 1024 / 1024
+
+
+def download_models(models):
+    for model in tqdm(models, desc="Downloading models"):
+        logger.info(f"Downloading {model}")
+        for i in range(0, 5):
+            curr_dir = f"{model}/train_{i}/best_model/"
+            os.makedirs(curr_dir, exist_ok=True)
+            os.system(
+                f"wget -q https://huggingface.co/researchaccount/{model}/resolve/main/train_{i}/best_model/config.json -P {curr_dir}"
+            )
+            os.system(
+                f"wget -q https://huggingface.co/researchaccount/{model}/resolve/main/train_{i}/best_model/pytorch_model.bin -P {curr_dir}"
+            )
+            os.system(
+                f"wget -q https://huggingface.co/researchaccount/{model}/resolve/main/train_{i}/best_model/special_tokens_map.json -P {curr_dir}"
+            )
+            os.system(
+                f"wget -q https://huggingface.co/researchaccount/{model}/resolve/main/train_{i}/best_model/tokenizer_config.json -P {curr_dir}"
+            )
+            os.system(
+                f"wget -q https://huggingface.co/researchaccount/{model}/resolve/main/train_{i}/best_model/training_args.bin -P {curr_dir}"
+            )
+            os.system(
+                f"wget -q https://huggingface.co/researchaccount/{model}/resolve/main/train_{i}/best_model/vocab.txt -P {curr_dir}"
+            )
+
+
+def softmax(x):
+    return np.exp(x) / sum(np.exp(x))
+
+
+def ga(file):
+    code = """
+    <!-- Global site tag (gtag.js) - Google Analytics -->
+        <script async src="https://www.googletagmanager.com/gtag/js?id=G-NH9HWCW08F"></script>
+        <script>
+        window.dataLayer = window.dataLayer || [];
+        function gtag(){dataLayer.push(arguments);}
+        gtag('js', new Date());
+        gtag('config', 'G-NH9HWCW08F');
+    </script>
+    """
+
+    a = os.path.dirname(file) + "/static/index.html"
+    with open(a, "r") as f:
+        data = f.read()
+        if len(re.findall("G-", data)) == 0:
+            with open(a, "w") as ff:
+                newdata = re.sub("<head>", "<head>" + code, data)
+                ff.write(newdata)

packages.txt

@@ -0,0 +1,2 @@
+openjdk-11-jre
+curl

requirements.txt

@@ -0,0 +1,17 @@
+streamlit==0.84.2
+arabic-reshaper==2.1.3
+python-bidi==0.4.2
+PyArabic
+farasapy==0.0.14
+emoji==1.4.2
+awesome_streamlit
+torch==1.9.0
+transformers==4.10.0
+psutil==5.8.0
+fuzzysearch==0.7.3
+more-itertools==8.9.0
+cookiecutter
+git+https://github.com/dantru7/Google-Search-API
+codetiming==1.3.0
+htbuilder
+wikipedia==1.4.0

test.py

@@ -0,0 +1,10 @@
+#%%
+from transformers import GPT2Tokenizer
+
+# %%
+tok = GPT2Tokenizer.from_pretrained("D:/ML/Models/aragpt2-medium", use_fast=False)
+# %%
+tok.pad_token = tok.eos_token
+#%%
+tok.pad_token_id = [tok.eos_token_id]
+# %%