refactor app, prepare for second prediction

.gitignore

@@ -0,0 +1,5 @@
+venv/
+llama*
+.DS_Store
+.vscode/
+old/

app.py

@@ -3,84 +3,147 @@
 
 import gradio as gr
 import numpy as np
-import requests
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, TextClassificationPipeline, pipeline
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    TextClassificationPipeline,
+    pipeline,
+)
 from langdetect import detect
 from matplotlib import pyplot as plt
 import imageio
 
-# Load the model
-model = AutoModelForSequenceClassification.from_pretrained("saved_model") 
-tokenizer = AutoTokenizer.from_pretrained("saved_model")
-pipe = TextClassificationPipeline(model=model, tokenizer=tokenizer)
+# move constants into extra file
+ML_MODEL_SURE = 0.6
+UNKNOWN_LANG_TEXT = (
+    "The language is not recognized, it must be either in German or in French."
+)
+PLACEHOLDER_TEXT = "Geben Sie bitte den Titel und den Sumbmitted Text des Vorstoss ein.\nVeuillez entrer le titre et le Submitted Text de la requête."
 
-# Function called by the UI
-def attribution(text):
-    
-    # Clean the plot
-    plt.clf()
-    
-    # Detect the language
-    language = detect(text)
-     
-    # Translate the input in german if necessary
-    if language == 'fr':
-        translator = pipeline("translation", model="Helsinki-NLP/opus-mt-fr-de")
-        translatedText = translator(text[0:1000])
-        text = translatedText[0]["translation_text"]
-    elif language != 'de':
-        return "The language is not recognized, it must be either in German or in French.", None
-    
-    # Set the bars of the bar chart
-    bars = ""
-    if language == 'fr':
-        bars = ("DDPS", "DFI", "AS-MPC", "DFJP", "DEFR", "DETEC", "DFAE", "Parl", "ChF", "DFF", "AF", "TF")
-    else:
-        bars = ("VBS", "EDI", "AB-BA", "EJPD", "WBF", "UVEK", "EDA", "Parl", "BK", "EFD", "BV", "BGer")
-
-    # Make the prediction with the 1000 first characters 
-    results = pipe(text[0:1000], return_all_scores=True)
-    rates = [row["score"] for row in results[0]]
-    
-    # Bar chart
-    y_pos = np.arange(len(bars))
+UNSURE_DE_TEXT = "Das ML-Modell ist nicht sicher. Das Departement könnte sein : \n\n"
+UNSURE_FR_TEXT = "Le modèle ML n'est pas sûr. Le département pourrait être : \n\n"
+BARS_DEP_FR = (
+    "DDPS",
+    "DFI",
+    "AS-MPC",
+    "DFJP",
+    "DEFR",
+    "DETEC",
+    "DFAE",
+    "Parl",
+    "ChF",
+    "DFF",
+    "AF",
+    "TF",
+)
+BARS_DEP_DE = (
+    "VBS",
+    "EDI",
+    "AB-BA",
+    "EJPD",
+    "WBF",
+    "UVEK",
+    "EDA",
+    "Parl",
+    "BK",
+    "EFD",
+    "BV",
+    "BGer",
+)
+
+
+def load_model(modelFolder):
+    """Loads model from model_folder & creates a text classification pipeline."""
+    model = AutoModelForSequenceClassification.from_pretrained(modelFolder)
+    tokenizer = AutoTokenizer.from_pretrained(modelFolder)
+    pipe = TextClassificationPipeline(model=model, tokenizer=tokenizer)
+    return pipe
+
+
+def translate_to_de(inputText):
+    """Translates french user input to German for the model to reach better classification."""
+    translator = pipeline("translation", model="Helsinki-NLP/opus-mt-fr-de")
+    translatedText = translator(inputText[0:1000])
+    text = translatedText[0]["translation_text"]
+    return text
+
+
+def create_bar_plot(rates, language):
+    barnames = BARS_DEP_FR if language == "fr" else BARS_DEP_DE
+
+    y_pos = np.arange(len(barnames))
     plt.barh(y_pos, rates)
-    plt.yticks(y_pos, bars)
-    
-    # Set the output text
-    name = ""    
+    plt.yticks(y_pos, barnames)
+
+    # Save the bar chart as png and load it (enables better display)
+    plt.savefig("rates.png")
+    im = imageio.v2.imread("rates.png")
+
+    return im, barnames
+
+
+def show_chosen_category(barnames, rates, language):
+    """Creates the output text
+    - adds disclaimer if ML model is not sure
+    - when unsure, adds all categories with prob. > 10% to output"""
     maxRate = np.max(rates)
     maxIndex = np.argmax(rates)
 
+    distance = "\t\t\t\t\t"
+
     # ML model not sure if highest probability < 60%
-    if maxRate < 0.6:
-        # de / fr
-        if language == 'de':
-            name = "Das ML-Modell ist nicht sicher. Das Departement könnte sein : \n\n"
-        else:
-            name = "Le modèle ML n'est pas sûr. Le département pourrait être : \n\n"
-        i = 0
+    if maxRate < ML_MODEL_SURE:
+        name = UNSURE_FR_TEXT if language == "fr" else UNSURE_DE_TEXT
+
         # Show each department that has a probability > 10%
+        i = 0
         while i == 0:
             if rates[maxIndex] >= 0.1:
-                name = name + "\t" + str(rates[maxIndex])[2:4] + "%" + "\t\t\t\t\t" + bars[maxIndex] + "\n"
+                chosenScore = str(rates[maxIndex])[2:4]
+                chosenCat = barnames[maxIndex]
+                name = name + "\t" + chosenScore + "%" + distance + chosenCat + "\n"
                 rates[maxIndex] = 0
                 maxIndex = np.argmax(rates)
             else:
                 i = 1
+
     # ML model pretty sure, show only one department
     else:
-        name =  str(maxRate)[2:4] + "%" + "\t\t\t\t\t\t" + bars[maxIndex]
-    
-    # Save the bar chart as png and load it (enables better display)
-    plt.savefig('rates.png')
-    im = imageio.imread('rates.png')
+        name = str(maxRate)[2:4] + "%" + distance + barnames[maxIndex]
+
+    return name
+
+
+pipeDep = load_model("saved_model_dep")
+# pipeOffice = load_model("saved_model_office")
+
+
+# Function called by the UI
+def attribution(inputText):
+    plt.clf()
+    language = detect(inputText)
+
+    # Translate the input to german if necessary
+    if language == "fr":
+        inputText = translate_to_de(inputText)
+    elif language != "de":
+        return UNKNOWN_LANG_TEXT, None
+
+    # Make the prediction with the 1000 first characters
+    prediction = pipeDep(inputText[0:1000], return_all_scores=True)
+    rates = [row["score"] for row in prediction[0]]
+
+    # Create barplot & output text
+    im, barnames = create_bar_plot(rates, language)
+    chosenCategoryText = show_chosen_category(barnames, rates, language)
+
+    return chosenCategoryText, im
 
-    return name, im
- 
 
 # display the UI
-interface = gr.Interface(fn=attribution, 
-                         inputs=[gr.inputs.Textbox(lines=20, placeholder="Geben Sie bitte den Titel und den Sumbmitted Text des Vorstoss ein.\nVeuillez entrer le titre et le Submitted Text de la requête.")],
-                          outputs=['text', 'image'])
-interface.launch()
+interface = gr.Interface(
+    fn=attribution,
+    inputs=[gr.components.Textbox(lines=20, placeholder=PLACEHOLDER_TEXT)],
+    outputs=["text", "image"],
+)
+interface.launch()

requirements.txt

@@ -4,4 +4,8 @@ langdetect
 matplotlib
 imageio
 torch
-sentencepiece
+sentencepiece
+
+gradio
+langdetect
+imageio

saved_model/config.json

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:f1dd5122dedc8fdf6eb1ec32b25f3246f8c3c64432abfd4d9bad4b626f378fc4
-size 1255

saved_model/special_tokens_map.json

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:303df45a03609e4ead04bc3dc1536d0ab19b5358db685b6f3da123d05ec200e3
-size 112

saved_model/tokenizer.json

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5d6f6affc6b91020cabef56fe9289907e34a89e7f3463a93250c0d94cc61000d
-size 726371

saved_model/tokenizer_config.json

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5ed472c8edcb18869d09d7bc852465911b105dd301fda14b4283b01577a5ebd7
-size 327

saved_model/vocab.txt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:982f8396ec746db0ed414dcc4789398ab6b365663cada50f776afb905dacbb61
-size 254729

saved_model_dep/config.json

@@ -0,0 +1,54 @@
+{
+  "_name_or_path": "bert-base-german-cased",
+  "architectures": [
+    "BertForSequenceClassification"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1",
+    "2": "LABEL_2",
+    "3": "LABEL_3",
+    "4": "LABEL_4",
+    "5": "LABEL_5",
+    "6": "LABEL_6",
+    "7": "LABEL_7",
+    "8": "LABEL_8",
+    "9": "LABEL_9",
+    "10": "LABEL_10",
+    "11": "LABEL_11"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1,
+    "LABEL_10": 10,
+    "LABEL_11": 11,
+    "LABEL_2": 2,
+    "LABEL_3": 3,
+    "LABEL_4": 4,
+    "LABEL_5": 5,
+    "LABEL_6": 6,
+    "LABEL_7": 7,
+    "LABEL_8": 8,
+    "LABEL_9": 9
+  },
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "problem_type": "single_label_classification",
+  "torch_dtype": "float32",
+  "transformers_version": "4.19.2",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30000
+}

saved_model_dep/special_tokens_map.json

@@ -0,0 +1 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

saved_model_dep/tokenizer_config.json

@@ -0,0 +1 @@
+{"do_lower_case": false, "unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]", "tokenize_chinese_chars": true, "strip_accents": null, "model_max_length": 512, "special_tokens_map_file": null, "name_or_path": "bert-base-german-cased", "tokenizer_class": "BertTokenizer"}