app.py

import gradio as gr
import numpy as np
import tensorflow as tf
from nltk import ngrams as NGrams
from transformers import TFBertForTokenClassification, AutoTokenizer
import pickle
import spacy
from spacy.tokens import Doc, Span
from spacy import displacy


with open("tokenizer.json") as file:
    data = file.read()
    tokenizer = tf.keras.preprocessing.text.tokenizer_from_json(data)


model_ = TFBertForTokenClassification.from_pretrained(
    "NER_Darija_NGrams_BERT_mode"
)

labels = np.array(
    ["O", "B-LOC", "B-MISC", "B-ORG", "B-PER", "I-LOC", "I-MISC", "I-ORG", "I-PER"]
)
label_mapping = {label: index for index, label in enumerate(labels)}
index2label = {v: k for k, v in label_mapping.items()}


def predict_ner(text):
    darija_tokens = text.split()
    n = 16  # Size of the n-grams
    max_length = 16

    # Padding and creating n-grams
    padded_tokens = (
        darija_tokens + ["<PAD>"] * (n - len(darija_tokens))
        if len(darija_tokens) < n
        else darija_tokens
    )
    test_ngrams = list(NGrams(padded_tokens, n))
    test_ngrams = [" ".join(ngram) for ngram in test_ngrams]

    # Model prediction
    X_darija = tokenizer.texts_to_sequences(test_ngrams)
    X_darija = tf.keras.preprocessing.sequence.pad_sequences(
        X_darija, maxlen=max_length, padding="post", truncating="post"
    )
    pred = model_(X_darija)
    probabilities = pred["logits"]

    # Accumulate probabilities for each token
    token_tag_prob = {token: np.zeros(9) for token in darija_tokens}
    for ngram, prob in zip(test_ngrams, probabilities):
        ngram_tokens = ngram.split()
        for i, token in enumerate(ngram_tokens):
            if token in token_tag_prob:
                token_tag_prob[token] = np.maximum(token_tag_prob[token], prob[i])

    # Determine the final label for each token
    final_labels = {
        token: index2label[np.argmax(probs)]
        for token, probs in token_tag_prob.items()
        if token != "<PAD>"
    }

    # Create a blank spaCy nlp pipeline
    nlp = spacy.blank("ar")

    # Create a Doc Doc object
    doc = Doc(
        nlp.vocab, words=list(final_labels.keys()), ents=list(final_labels.values())
    )

    # Visualize using displacy
    html = displacy.render(doc, style="ent")

    return html


# Define your Gradio interface
iface = gr.Interface(
    fn=predict_ner,
    inputs=gr.Textbox(lines=2, placeholder="Enter a sentence in Darija..."),
    outputs="html",
    examples=[
        "فنفس لعام سنات رجاء  كونطرة مع الشركة ديال روطانا  وخرجات الديسك ديالها استغربت لحال الدنيا ونجحات ليها لأغنية فلعالم لعربي",
        "لمرنيسي  معروفة بلمقاربة السوسيوسياسية فاش كتبغي تدرس لجندر ولهوية لجنسانية وتركيزها كيكون بلمجمل على لمغريب  لحاجة لي خلاتها تعرف برا لبلاد هي كونها فيمينيست مسلمة كانت لمرنيسي  محاضرة فلافاك ديال محمد لخامس  ف الرباط  وباحتة ف لمعهد لجامعي ديال لبحت لعلمي  لي كاين فا لرباط  نيت",
        "الشركة د لبراسريات د لمغريب  كاينين ف فاس  طانجا ؤ كازا  ؤ كاين واحد لوحدة ديال التعمار د لقراعي ف مراكش"
    ],
)

iface.launch()