added freematch test

app.py

@@ -1,7 +1,80 @@
+import re
 import gradio as gr
+from transformers import AutoTokenizer
+from unidecode import unidecode
+from models import *
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+tok = AutoTokenizer.from_pretrained("readerbench/RoBERT-base")
+
+def preprocess(x):
+    """Preprocess input string x"""
+
+    s = unidecode(x)
+    s = str.lower(s)
+    s = re.sub(r"\[[a-z]+\]","", s)
+    s = re.sub(r"\*","", s)
+    s = re.sub(r"[^a-zA-Z0-9]+"," ",s)
+    s = re.sub(r" +"," ",s)
+    s = re.sub(r"(.)\1+",r"\1",s)
+
+    return s
+
+label_names = ["ABUSE", "INSULT", "OTHER", "PROFANITY"]
+
+def ssl_predict(in_text, model_type):
+    """main predict function"""
+
+    preprocessed = preprocess(in_text)
+    toks = tok(
+        preprocessed,
+        padding="max_length",
+        max_length=96,
+        truncation=True,
+        return_tensors="tf"
+    )
+    if model_type == "freematch":
+        model = FixMatchTune(encoder_name="andrei-saceleanu/ro-offense-freematch")
+        model.cls_head.load_weights("./checkpoints/freematch_tune")
+
+    preds, _ = model([toks["input_ids"],toks["attention_mask"]], training=False)
+
+    probs = list(preds[0].numpy())
+
+    return {k:v for k, v in zip(label_names, probs)}
+
+
+
+
+with gr.Blocks() as ssl_interface:
+    with gr.Row():
+        with gr.Column():
+            in_text = gr.Textbox(label="Input text")
+            model_list = gr.Dropdown(
+                choices=["fixmatch", "freematch", "mixmatch"],
+                max_choices=1,
+                label="Training method",
+                allow_custom_value=False,
+                info="Select trained model according to different SSL techniques from paper",
+            )
+
+            with gr.Row():
+                clear_btn = gr.Button(value="Clear")
+                submit_btn = gr.Button(value="Submit")
+
+        with gr.Column():
+            out_field = gr.Label(num_top_classes=4,label="Prediction")
+
+    submit_btn.click(
+        fn=ssl_predict,
+        inputs=[in_text, model_list],
+        outputs=[out_field]
+    )
+
+    clear_btn.click(
+        fn=lambda: [None for _ in range(2)],
+        inputs=None,
+        outputs=[in_text, out_field]
+    )
+
+ssl_interface.launch(server_name="0.0.0.0", server_port=7860)

models.py

@@ -0,0 +1,37 @@
+import tensorflow as tf
+from transformers import TFAutoModel
+
+
+class FixMatchTune(tf.keras.Model):
+    def __init__(
+        self,
+        encoder_name="readerbench/RoBERT-base",
+        num_classes=4,
+        **kwargs
+    ):
+        super(FixMatchTune,self).__init__(**kwargs)
+
+        self.bert = TFAutoModel.from_pretrained(encoder_name)
+        self.num_classes = num_classes
+        self.weak_augment = tf.keras.layers.GaussianNoise(stddev=0.5)
+        self.strong_augment = tf.keras.layers.GaussianNoise(stddev=5)
+
+        self.cls_head = tf.keras.Sequential([
+            tf.keras.layers.Dense(256,activation="relu"),
+            tf.keras.layers.Dropout(0.2),
+            tf.keras.layers.Dense(64,activation="relu"),
+            tf.keras.layers.Dense(self.num_classes, activation="softmax")
+        ])
+
+    def call(self, inputs, training):
+        ids, mask = inputs
+
+        embeds = self.bert(input_ids=ids, attention_mask=mask,training=training).pooler_output
+
+        strongs = self.strong_augment(embeds,training=training)
+        weaks = self.weak_augment(embeds,training=training)
+
+        strong_preds = self.cls_head(strongs,training=training)
+        weak_preds = self.cls_head(weaks,training=training)
+
+        return weak_preds, strong_preds