Added model selection feature

app.py

@@ -1,6 +1,5 @@
 import logging
 import os
-import tempfile
 from pathlib import Path
 from typing import List, Tuple
 
@@ -8,6 +7,7 @@ import gradio as gr
 import pandas as pd
 import spacy
 import torch
+from dante_tokenizer import DanteTokenizer
 from transformers import AutoModelForTokenClassification, AutoTokenizer
 
 from preprocessing import expand_contractions
@@ -17,27 +17,56 @@ try:
 except Exception:
     os.system("python -m spacy download pt_core_news_sm")
     nlp = spacy.load("pt_core_news_sm")
-
-model = AutoModelForTokenClassification.from_pretrained("Emanuel/porttagger-news-base")
-tokenizer = AutoTokenizer.from_pretrained("Emanuel/porttagger-news-base")
+dt_tokenizer = DanteTokenizer()
+
+default_model = "News"
+model_choices = {
+    "News": "Emanuel/porttagger-news-base",
+    "Tweets": "Emanuel/porttagger-tweets-base",
+    "Oil and Gas": "Emanuel/porttagger-oilgas-base",
+    "Multigenre": "Emanuel/porttagger-base",
+}
+pre_tokenizers = {
+    "News": nlp,
+    "Tweets": dt_tokenizer.tokenize,
+    "Oil and Gas": nlp,
+    "Multigenre": nlp,
+}
 logger = logging.getLogger()
 logger.setLevel(logging.DEBUG)
 
+class MyApp:
+    def __init__(self) -> None:
+        self.model = None
+        self.tokenizer = None
+        self.pre_tokenizer = None
+        self.load_model()
+
+    def load_model(self, model_name: str = default_model):
+        if model_name not in model_choices.keys():
+            logger.error("Selected model is not supported, resetting to the default model.")
+            model_name = default_model
+        self.model = AutoModelForTokenClassification.from_pretrained(model_choices[model_name])
+        self.tokenizer = AutoTokenizer.from_pretrained(model_choices[model_name])
+        self.pre_tokenizer = pre_tokenizers[model_name]
+
+myapp = MyApp()
 
-def predict(text, nlp, logger=None) -> Tuple[List[str], List[str]]:
-    doc = nlp(text)
-    tokens = [token.text for token in doc]
+def predict(text, logger=None) -> Tuple[List[str], List[str]]:
+    doc = myapp.pre_tokenizer(text)
+    tokens = [token.text if not isinstance(token, str) else token for token in doc]
 
     logger.info("Starting predictions for sentence: {}".format(text))
+    print("Using model {}".format(myapp.model.config.__dict__["_name_or_path"]))
 
-    input_tokens = tokenizer(
+    input_tokens = myapp.tokenizer(
         tokens,
         return_tensors="pt",
         is_split_into_words=True,
         return_offsets_mapping=True,
         return_special_tokens_mask=True,
     )
-    output = model(input_tokens["input_ids"])
+    output = myapp.model(input_tokens["input_ids"])
 
     i_token = 0
     labels = []
@@ -49,7 +78,7 @@ def predict(text, nlp, logger=None) -> Tuple[List[str], List[str]]:
     ):
         if is_special_token or off[0] > 0:
             continue
-        label = model.config.__dict__["id2label"][int(pred.argmax(axis=-1))]
+        label = myapp.model.config.__dict__["id2label"][int(pred.argmax(axis=-1))]
         if logger is not None:
             logger.info("{}, {}, {}".format(off, tokens[i_token], label))
         labels.append(label)
@@ -63,7 +92,11 @@ def predict(text, nlp, logger=None) -> Tuple[List[str], List[str]]:
 
 def text_analysis(text):
     text = expand_contractions(text)
-    tokens, labels, scores = predict(text, nlp, logger)
+    tokens, labels, scores = predict(text, logger)
+    if len(labels) != len(tokens):
+        m = len(tokens) - len(labels)
+        labels += [None] * m
+        scores += [0] * m
     pos_count = pd.DataFrame(
         {
             "token": tokens,
@@ -99,7 +132,7 @@ def batch_analysis(input_file):
         sent = expand_contractions(sent)
         conllu_output.append("# sent_id = {}-{}\n".format(name, i + 1))
         conllu_output.append("# text = {}\n".format(sent))
-        tokens, labels, scores = predict(sent, nlp, logger)
+        tokens, labels, scores = predict(sent, logger)
         for j, (token, label) in enumerate(zip(tokens, labels)):
             conllu_output.append(
                 "{}\t{}\t_\t{}".format(j + 1, token, label) + "\t_" * 5 + "\n"
@@ -119,6 +152,8 @@ bottom_html = open("bottom.html").read()
 
 with gr.Blocks(css=css) as demo:
     gr.HTML(top_html)
+    select_model = gr.Dropdown(choices=list(model_choices.keys()), label="Tagger model", value=default_model)
+    select_model.change(myapp.load_model, inputs=[select_model])
     with gr.Tab("Single sentence"):
         text = gr.Textbox(placeholder="Enter your text here...", label="Input")
         examples = gr.Examples(
@@ -167,4 +202,4 @@ with gr.Blocks(css=css) as demo:
     gr.HTML(bottom_html)
 
 
-demo.launch(debug=True, server_port=15000)
+demo.launch(debug=True)