import streamlit as st

# Custom CSS for better styling
st.markdown("""
    <style>
        .main-title {
            font-size: 36px;
            color: #4A90E2;
            font-weight: bold;
            text-align: center;
        }
        .sub-title {
            font-size: 24px;
            color: #4A90E2;
            margin-top: 20px;
        }
        .section {
            background-color: #f9f9f9;
            padding: 15px;
            border-radius: 10px;
            margin-top: 20px;
        }
        .section h2 {
            font-size: 22px;
            color: #4A90E2;
        }
        .section p, .section ul {
            color: #666666;
        }
        .link {
            color: #4A90E2;
            text-decoration: none;
        }
        .benchmark-table {
            width: 100%;
            border-collapse: collapse;
            margin-top: 20px;
        }
        .benchmark-table th, .benchmark-table td {
            border: 1px solid #ddd;
            padding: 8px;
            text-align: left;
        }
        .benchmark-table th {
            background-color: #4A90E2;
            color: white;
        }
        .benchmark-table td {
            background-color: #f2f2f2;
        }
    </style>
""", unsafe_allow_html=True)

# Title
st.markdown('<div class="main-title">Introduction to ALBERT Annotators in Spark NLP</div>', unsafe_allow_html=True)

# Subtitle
st.markdown("""
<div class="section">
    <p>ALBERT (A Lite BERT) offers a more efficient alternative to BERT by implementing two parameter-reduction techniques: splitting the embedding matrix and using repeating layers. It maintains high performance while being more memory-efficient. Below, we provide an overview of the ALBERT annotator for token classification:</p>
</div>
""", unsafe_allow_html=True)

tab1, tab2, tab3 = st.tabs(["ALBERT for Token Classification", "ALBERT for Sequence Classification", "ALBERT for Question Answering"])

with tab1:
    st.markdown("""
    <div class="section">
        <h2>ALBERT for Token Classification</h2>
        <p>The <strong>AlbertForTokenClassification</strong> annotator is designed for Named Entity Recognition (NER) tasks using ALBERT. This model efficiently handles token classification, enabling the identification and classification of entities in text. The ALBERT model, with its parameter-reduction techniques, achieves state-of-the-art performance while being more lightweight compared to BERT.</p>
        <p>Token classification with ALBERT enables:</p>
        <ul>
            <li><strong>Named Entity Recognition (NER):</strong> Identifying and classifying entities such as names, organizations, locations, and other predefined categories.</li>
            <li><strong>Information Extraction:</strong> Extracting key information from unstructured text for further analysis.</li>
            <li><strong>Text Categorization:</strong> Enhancing document retrieval and categorization based on entity recognition.</li>
        </ul>
        <p>Here is an example of how ALBERT token classification works:</p>
        <table class="benchmark-table">
            <tr>
                <th>Entity</th>
                <th>Label</th>
            </tr>
            <tr>
                <td>Google</td>
                <td>ORG</td>
            </tr>
            <tr>
                <td>Satya Nadella</td>
                <td>PER</td>
            </tr>
            <tr>
                <td>Seattle</td>
                <td>LOC</td>
            </tr>
        </table>
    </div>
    """, unsafe_allow_html=True)

    # ALBERT Token Classification - NER CoNLL
    st.markdown('<div class="sub-title">ALBERT Token Classification - NER CoNLL</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <p>The <strong>albert_base_token_classifier_conll03</strong> is a fine-tuned ALBERT model for token classification tasks, specifically adapted for Named Entity Recognition (NER) on the CoNLL-03 dataset. It recognizes four types of entities: location (LOC), organizations (ORG), person (PER), and Miscellaneous (MISC).</p>
    </div>
    """, unsafe_allow_html=True)

    # How to Use the Model - Token Classification
    st.markdown('<div class="sub-title">How to Use the Model</div>', unsafe_allow_html=True)
    st.code('''
    from sparknlp.base import *
    from sparknlp.annotator import *
    from pyspark.ml import Pipeline
    from pyspark.sql.functions import col, expr

    document_assembler = DocumentAssembler() \\
        .setInputCol('text') \\
        .setOutputCol('document')

    tokenizer = Tokenizer() \\
        .setInputCols(['document']) \\
        .setOutputCol('token')

    tokenClassifier = AlbertForTokenClassification \\
        .pretrained('albert_base_token_classifier_conll03', 'en') \\
        .setInputCols(['token', 'document']) \\
        .setOutputCol('ner') \\
        .setCaseSensitive(True) \\
        .setMaxSentenceLength(512)

    # Convert NER labels to entities
    ner_converter = NerConverter() \\
        .setInputCols(['document', 'token', 'ner']) \\
        .setOutputCol('entities')

    pipeline = Pipeline(stages=[
        document_assembler,
        tokenizer,
        tokenClassifier,
        ner_converter
    ])

    example = spark.createDataFrame([["My name is John!"]]).toDF("text")
    result = pipeline.fit(example).transform(example)

    result.select(
        expr("explode(entities) as ner_chunk")
    ).select(
        col("ner_chunk.result").alias("chunk"),
        col("ner_chunk.metadata.entity").alias("ner_label")
    ).show(truncate=False)
    ''', language='python')

    # Results
    st.text("""
    +-----+---------+
    |chunk|ner_label|
    +-----+---------+
    |John |PER      |
    +-----+---------+
    """)

    # Performance Metrics
    st.markdown('<div class="sub-title">Performance Metrics</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <p>Here are the detailed performance metrics for the ALBERT token classification model:</p>
        <table class="benchmark-table">
            <tr>
                <th>Entity</th>
                <th>Precision</th>
                <th>Recall</th>
                <th>F1-Score</th>
                <th>Support</th>
            </tr>
            <tr>
                <td>B-LOC</td>
                <td>0.95</td>
                <td>0.97</td>
                <td>0.96</td>
                <td>1837</td>
            </tr>
            <tr>
                <td>B-MISC</td>
                <td>0.87</td>
                <td>0.86</td>
                <td>0.87</td>
                <td>922</td>
            </tr>
            <tr>
                <td>B-ORG</td>
                <td>0.90</td>
                <td>0.91</td>
                <td>0.90</td>
                <td>1341</td>
            </tr>
            <tr>
                <td>B-PER</td>
                <td>0.91</td>
                <td>0.97</td>
                <td>0.94</td>
                <td>1842</td>
            </tr>
            <tr>
                <td>I-LOC</td>
                <td>0.88</td>
                <td>0.86</td>
                <td>0.87</td>
                <td>257</td>
            </tr>
            <tr>
                <td>I-MISC</td>
                <td>0.78</td>
                <td>0.76</td>
                <td>0.77</td>
                <td>346</td>
            </tr>
            <tr>
                <td>I-ORG</td>
                <td>0.84</td>
                <td>0.85</td>
                <td>0.85</td>
                <td>751</td>
            </tr>
            <tr>
                <td>I-PER</td>
                <td>0.97</td>
                <td>0.92</td>
                <td>0.94</td>
                <td>1307</td>
            </tr>
            <tr>
                <td>O</td>
                <td>0.99</td>
                <td>0.99</td>
                <td>0.99</td>
                <td>42759</td>
            </tr>
            <tr>
                <td>average</td>
                <td>0.92</td>
                <td>0.92</td>
                <td>0.92</td>
                <td>52000</td>
            </tr>
        </table>
    </div>
                
    """, unsafe_allow_html=True)
    # Model Info Section
    st.markdown('<div class="sub-title">Model Info</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <ul>
            <li><strong>Model Name:</strong> ALBERT for Token Classification</li>
            <li><strong>Pretrained Model:</strong> albert_base_token_classifier_conll03</li>
            <li><strong>Training Dataset:</strong> CoNLL-03</li>
            <li><strong>Languages Supported:</strong> English</li>
            <li><strong>Use Cases:</strong>
                <ul>
                    <li>Named Entity Recognition (NER)</li>
                    <li>Information Extraction</li>
                    <li>Text Categorization</li>
                </ul>
            </li>
            <li><strong>Performance:</strong> High accuracy with a focus on memory efficiency</li>
            <li><strong>Implementation:</strong> Spark NLP</li>
            <li><strong>Resource Requirements:</strong> Moderate computational resources; suitable for production environments with optimization</li>
        </ul>
    </div>
    """, unsafe_allow_html=True)

    # References Section
    st.markdown('<div class="sub-title">References</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <ul>
            <li><a class="link" href="https://arxiv.org/abs/1909.11942" target="_blank">Lan, Z., Chen, J., Goodman, S., Gimpel, K., Sharma, P., & Soricut, R. (2019). ALBERT: A Lite BERT for Self-supervised Learning of Language Representations. arXiv preprint arXiv:1909.11942.</a></li>
            <li><a class="link" href="https://github.com/google-research/albert" target="_blank">Google Research's ALBERT GitHub Repository</a></li>
            <li><a class="link" href="https://sparknlp.org/2022/06/15/albert_base_qa_squad2_en_3_0.html" target="_blank">Spark NLP Model - albert_base_qa_squad2</a></li>
            <li><a class="link" href="https://nlp.stanford.edu/projects/conll2003/" target="_blank">CoNLL-03 Named Entity Recognition Dataset</a></li>
        </ul>
    </div>
    """, unsafe_allow_html=True)

with tab2:
    st.markdown("""
    <div class="section">
        <h2>ALBERT for Sequence Classification</h2>
        <p>The <strong>AlbertForSequenceClassification</strong> annotator is tailored for tasks like sentiment analysis or multi-class text classification using the ALBERT model. This model efficiently handles sequence classification, achieving state-of-the-art performance with reduced parameters compared to BERT.</p>
        <p>Sequence classification with ALBERT enables:</p>
        <ul>
            <li><strong>Sentiment Analysis:</strong> Determining the sentiment expressed in text, such as positive, negative, or neutral.</li>
            <li><strong>Multi-Class Text Classification:</strong> Categorizing text into predefined classes, such as news categories or topics.</li>
            <li><strong>Document Classification:</strong> Enhancing search and categorization of documents based on content classification.</li>
        </ul>
        <p>Here is an example of how ALBERT sequence classification works:</p>
        <table class="benchmark-table">
            <tr>
                <th>Text</th>
                <th>Label</th>
            </tr>
            <tr>
                <td>Disney Comics was a comic book publishing company operated by The Walt Disney Company which ran from 1990 to 1993.</td>
                <td>Business</td>
            </tr>
        </table>
    </div>
    """, unsafe_allow_html=True)

    # ALBERT Sequence Classification - AG News
    st.markdown('<div class="sub-title">ALBERT Sequence Classification - AG News</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <p>The <strong>albert_base_sequence_classifier_ag_news</strong> is a fine-tuned ALBERT model for sequence classification tasks, specifically adapted for text classification on the AG News dataset. It recognizes four categories: Business, Sci/Tech, Sports, and World.</p>
    </div>
    """, unsafe_allow_html=True)

    # How to Use the Model - Sequence Classification
    st.markdown('<div class="sub-title">How to Use the Model</div>', unsafe_allow_html=True)
    st.code('''
    from sparknlp.base import *
    from sparknlp.annotator import *
    from pyspark.ml import Pipeline
    from pyspark.sql.functions import col, expr

    document_assembler = DocumentAssembler() \\
        .setInputCol('text') \\
        .setOutputCol('document')

    tokenizer = Tokenizer() \\
        .setInputCols(['document']) \\
        .setOutputCol('token')

    sequenceClassifier = AlbertForSequenceClassification \\
        .pretrained('albert_base_sequence_classifier_ag_news', 'en') \\
        .setInputCols(['token', 'document']) \\
        .setOutputCol('class') \\
        .setCaseSensitive(False) \\
        .setMaxSentenceLength(512)

    pipeline = Pipeline(stages=[
        document_assembler,
        tokenizer,
        sequenceClassifier    
    ])

    example = spark.createDataFrame([["Disney Comics was a comic book publishing company operated by The Walt Disney Company which ran from 1990 to 1993."]]).toDF("text")
    result = pipeline.fit(example).transform(example)

    result.select(
        expr("explode(class) as classification_result")
    ).select(
        col("classification_result.result").alias("category")
    ).show(truncate=False)
    ''', language='python')

    # Results
    st.text("""
    +---------+
    |category |
    +---------+
    |Business |
    +---------+
    """)

    # Performance Metrics
    st.markdown('<div class="sub-title">Performance Metrics</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <p>Here are the detailed performance metrics for the ALBERT sequence classification model on the AG News dataset:</p>
        <table class="benchmark-table">
            <tr>
                <th>Metric</th>
                <th>Score</th>
            </tr>
            <tr>
                <td>Accuracy</td>
                <td>0.9472</td>
            </tr>
            <tr>
                <td>F1-Score</td>
                <td>0.9472</td>
            </tr>
            <tr>
                <td>Precision</td>
                <td>0.9472</td>
            </tr>
            <tr>
                <td>Recall</td>
                <td>0.9472</td>
            </tr>
            <tr>
                <td>Evaluation Loss</td>
                <td>0.1882</td>
            </tr>
        </table>
    </div>
                
    """, unsafe_allow_html=True)
    
    # Model Info Section
    st.markdown('<div class="sub-title">Model Info</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <ul>
            <li><strong>Model Name:</strong> ALBERT for Sequence Classification</li>
            <li><strong>Pretrained Model:</strong> albert_base_sequence_classifier_ag_news</li>
            <li><strong>Training Dataset:</strong> AG News</li>
            <li><strong>Languages Supported:</strong> English</li>
            <li><strong>Use Cases:</strong>
                <ul>
                    <li>Sentiment Analysis</li>
                    <li>Multi-Class Text Classification</li>
                    <li>Document Classification</li>
                </ul>
            </li>
            <li><strong>Performance:</strong> High accuracy with a focus on memory efficiency</li>
            <li><strong>Implementation:</strong> Spark NLP</li>
            <li><strong>Resource Requirements:</strong> Moderate computational resources; suitable for production environments with optimization</li>
        </ul>
    </div>
    """, unsafe_allow_html=True)

    # References Section
    st.markdown('<div class="sub-title">References</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <ul>
            <li><a class="link" href="https://arxiv.org/abs/1909.11942" target="_blank">Lan, Z., Chen, J., Goodman, S., Gimpel, K., Sharma, P., & Soricut, R. (2019). ALBERT: A Lite BERT for Self-supervised Learning of Language Representations. arXiv preprint arXiv:1909.11942.</a></li>
            <li><a class="link" href="https://github.com/google-research/albert" target="_blank">Google Research's ALBERT GitHub Repository</a></li>
            <li><a class="link" href="https://sparknlp.org/2021/12/16/albert_base_sequence_classifier_ag_news_en.html" target="_blank">Spark NLP Model - albert_base_sequence_classifier_ag_news</a></li>
            <li><a class="link" href="https://huggingface.co/datasets/ag_news" target="_blank">AG News Dataset</a></li>
        </ul>
    </div>
    """, unsafe_allow_html=True)

with tab3:
    st.markdown("""
    <div class="section">
        <h2>ALBERT for Question Answering</h2>
        <p>The <strong>AlbertForQuestionAnswering</strong> annotator is specialized for tasks involving Question Answering (QA) using the ALBERT model. This model efficiently processes question-context pairs to provide accurate answers, making it ideal for QA systems and information retrieval applications.</p>
        <p>Question Answering with ALBERT enables:</p>
        <ul>
            <li><strong>Information Retrieval:</strong> Extracting precise answers from large text corpora based on user queries.</li>
            <li><strong>Knowledge Management:</strong> Enhancing customer support and information systems by providing accurate answers.</li>
            <li><strong>Contextual Understanding:</strong> Leveraging ALBERT’s capabilities to understand the context of questions and provide relevant answers.</li>
        </ul>
        <p>Here is an example of how ALBERT question answering works:</p>
        <table class="benchmark-table">
            <tr>
                <th>Question</th>
                <th>Context</th>
                <th>Answer</th>
            </tr>
            <tr>
                <td>What is my name?</td>
                <td>My name is Clara and I live in Berkeley.</td>
                <td>Clara</td>
            </tr>
        </table>
    </div>
    """, unsafe_allow_html=True)

    # ALBERT Question Answering - SQuAD2
    st.markdown('<div class="sub-title">ALBERT Question Answering - SQuAD2</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <p>The <strong>albert_base_qa_squad2</strong> is a fine-tuned ALBERT model for Question Answering tasks, specifically adapted for the SQuAD2 dataset. It is capable of answering questions based on the provided context with high accuracy.</p>
    </div>
    """, unsafe_allow_html=True)

    # How to Use the Model - Question Answering
    st.markdown('<div class="sub-title">How to Use the Model</div>', unsafe_allow_html=True)
    st.code('''
    from sparknlp.base import *
    from sparknlp.annotator import *
    from pyspark.ml import Pipeline
    from pyspark.sql.functions import col, expr

    documentAssembler = MultiDocumentAssembler() \\
        .setInputCols(["question", "context"]) \\
        .setOutputCols(["document_question", "document_context"])

    spanClassifier = AlbertForQuestionAnswering.pretrained("albert_base_qa_squad2","en") \\
        .setInputCols(["document_question", "document_context"]) \\
        .setOutputCol("answer") \\
        .setCaseSensitive(False)

    pipeline = Pipeline(stages=[documentAssembler, spanClassifier])

    data = spark.createDataFrame([["What is my name?", "My name is Clara and I live in Berkeley."]]).toDF("question", "context")
    result = pipeline.fit(data).transform(data)

    result.select(
        col("answer.result").alias("predicted_answer")
    ).show(truncate=False)
    ''', language='python')

    # Results
    st.text("""
    +----------------+
    |predicted_answer|
    +----------------+
    |[clara]         |
    +----------------+
    """)

    # Performance Metrics
    st.markdown('<div class="sub-title">Performance Metrics</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <p>The performance metrics of the ALBERT question answering model on a development subset of the SQuAD2 dataset are:</p>
        <table class="benchmark-table">
            <tr>
                <th>Metric</th>
                <th>Score</th>
            </tr>
            <tr>
                <td>Exact Match</td>
                <td>78.71%</td>
            </tr>
            <tr>
                <td>F1 Score</td>
                <td>81.89%</td>
            </tr>
            <tr>
                <td>Total</td>
                <td>6078</td>
            </tr>
            <tr>
                <td>HasAns Exact Match</td>
                <td>75.40%</td>
            </tr>
            <tr>
                <td>HasAns F1 Score</td>
                <td>82.04%</td>
            </tr>
            <tr>
                <td>HasAns Total</td>
                <td>2910</td>
            </tr>
            <tr>
                <td>NoAns Exact Match</td>
                <td>81.76%</td>
            </tr>
            <tr>
                <td>NoAns F1 Score</td>
                <td>81.76%</td>
            </tr>
            <tr>
                <td>NoAns Total</td>
                <td>3168</td>
            </tr>
            <tr>
                <td>Best Exact Match</td>
                <td>78.73%</td>
            </tr>
            <tr>
                <td>Best F1 Score</td>
                <td>81.91%</td>
            </tr>
        </table>
    </div>
    """, unsafe_allow_html=True)
    
    # Model Info Section
    st.markdown('<div class="sub-title">Model Info</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <ul>
            <li><strong>Model Name:</strong> ALBERT for Question Answering</li>
            <li><strong>Pretrained Model:</strong> albert_base_qa_squad2</li>
            <li><strong>Training Dataset:</strong> SQuAD2</li>
            <li><strong>Languages Supported:</strong> English</li>
            <li><strong>Use Cases:</strong>
                <ul>
                    <li>Information Retrieval</li>
                    <li>Knowledge Management</li>
                    <li>Contextual Understanding</li>
                </ul>
            </li>
            <li><strong>Performance:</strong> High accuracy with optimized resource usage</li>
            <li><strong>Implementation:</strong> Spark NLP</li>
            <li><strong>Resource Requirements:</strong> Moderate computational resources; suitable for production environments</li>
        </ul>
    </div>
    """, unsafe_allow_html=True)

    # References Section
    st.markdown('<div class="sub-title">References</div>', unsafe_allow_html=True)
    st.markdown("""
    <div class="section">
        <ul>
            <li><a class="link" href="https://arxiv.org/abs/1909.11942" target="_blank">Lan, Z., Chen, J., Goodman, S., Gimpel, K., Sharma, P., & Soricut, R. (2019). ALBERT: A Lite BERT for Self-supervised Learning of Language Representations. arXiv preprint arXiv:1909.11942.</a></li>
            <li><a class="link" href="https://sparknlp.org/2022/06/15/albert_base_qa_squad2_en_3_0.html" target="_blank">Spark NLP Model - albert_base_qa_squad2</a></li>
        </ul>
    </div>
    """, unsafe_allow_html=True)

# Community & Support
st.markdown('<div class="sub-title">Community & Support</div>', unsafe_allow_html=True)
st.markdown("""
<div class="section">
    <ul>
        <li><a class="link" href="https://sparknlp.org/" target="_blank">Official Website</a>: Documentation and examples</li>
        <li><a class="link" href="https://join.slack.com/t/spark-nlp/shared_invite/zt-198dipu77-L3UWNe_AJ8xqDk0ivmih5Q" target="_blank">Slack</a>: Live discussion with the community and team</li>
        <li><a class="link" href="https://github.com/JohnSnowLabs/spark-nlp" target="_blank">GitHub</a>: Bug reports, feature requests, and contributions</li>
    </ul>
</div>
""", unsafe_allow_html=True)