FEAT: Improved diagnostic mode with better output matrix

app.py

@@ -1,4 +1,6 @@
 import os
+import re
+import numpy as np
 import glob
 import json
 from typing import Dict, List, Tuple, Union
@@ -8,10 +10,9 @@ import pandas
 import streamlit as st
 import matplotlib.pyplot as plt
 
-
 from inference_tokenizer import NextSentencePredictionTokenizer
 from models import get_class
-from models import OwnBertForNextSentencePrediction
+
 
 def get_model(_model_path):
     print(f"Getting model at {_model_path}")
@@ -57,7 +58,6 @@ for model_path in models_path:
         model_data["path"] = model_path.replace("info.json", "")
         models[model_data["model"]] = model_data
 
-
 model_name = st.selectbox('Which model do you want to use?',
                           (x for x in sorted(models.keys())),
                           index=0)
@@ -78,15 +78,28 @@ def get_evaluation_data_from_json(_context: List) -> List[Tuple[List, str, str]]
     return output_data
 
 
-def get_evaluation_data_from_dialogue(_context: List) -> List[Tuple[List, str, Union[str, None]]]:
+control_sequence_regex_1 = re.compile(r"#.*? ")
+control_sequence_regex_2 = re.compile(r"#.*?\n")
+
+
+def _clean_conversational_line(_line: str):
+    _line = _line.replace("Bot: ", "")
+    _line = _line.replace("User: ", "")
+    _line = control_sequence_regex_1.sub("", _line)
+    _line = control_sequence_regex_2.sub("\n", _line)
+    return _line.strip()
+
+
+def get_evaluation_data_from_dialogue(_context: List[str]) -> List[Dict]:
     output_data = []
+    _context = list(map(lambda x: x.strip(), _context))
+    _context = list(filter(lambda x: len(x), _context))
     for idx, _line in enumerate(_context):
-        if idx == 0:
-            continue
         actual_context = _context[max(0, idx - 5):idx]
-        actual_sentence = _line
+        gradual_context_dict = {_line: []}
         for context_idx in range(len(actual_context)):
-            output_data.append((actual_context[-context_idx:], actual_sentence, None))
+            gradual_context_dict[_line].append(actual_context[-context_idx:])
+        output_data.append(gradual_context_dict)
     return output_data
 
 
@@ -97,7 +110,6 @@ option = st.selectbox("Choose type of input:",
                        "04 - JSON (example Elysai)",
                        "05 - Diagnostic mode"])
 
-
 with st.form("input_text"):
     if "01" in option:
         context = st.text_area("Insert context here (one turn per line):")
@@ -119,6 +131,7 @@ with st.form("input_text"):
             st.pyplot(fig)
     elif "02" in option or "03" in option or "04" in option:
         from data.example_data import ca_ood, elysai
+
         choices = [ca_ood, elysai]
         option: str
         # > Python 3.10
@@ -163,20 +176,27 @@ with st.form("input_text"):
                                                     "Probability (not-follow)"])
             st.dataframe(df)
     elif "05" in option:
+        context_size = 5
         context = st.text_area("Insert dialogue here (one turn per line):")
         submitted = st.form_submit_button("Submit")
         if submitted:
-            aggregated_result = []
-            data_for_evaluation = get_evaluation_data_from_dialogue(context.split("\n"))
+            data_for_evaluation = get_evaluation_data_from_dialogue(_clean_conversational_line(context).split("\n"))
+            lines = []
+            scores = np.zeros(shape=(len(data_for_evaluation), context_size))
             for datapoint in data_for_evaluation:
-                c, s, _ = datapoint
-                input_tensor = inference_tokenizer.get_item(context=c, actual_sentence=s)
-                output_model = model(**input_tensor.data).logits
-                output_model = torch.softmax(output_model, dim=-1).detach().numpy()[0]
-                prop_follow = output_model[0]
-                prop_not_follow = output_model[1]
+                for actual_sentence, contexts in datapoint.items():
+                    lines.append(actual_sentence)
+                    for c in contexts:
+                        input_tensor = inference_tokenizer.get_item(context=c, actual_sentence=actual_sentence)
+                        output_model = model(**input_tensor.data).logits
+                        output_model = torch.softmax(output_model, dim=-1).detach().numpy()[0]
+                        prop_follow = output_model[0]
+                        prop_not_follow = output_model[1]
+                        scores[len(lines) - 1][len(c) - 1] = prop_follow
 
-                aggregated_result.append((c, s, prop_follow))
+            aggregated_result = []
+            for idx, line in enumerate(lines):
+                aggregated_result.append([line] + scores[idx].tolist())
             st.table(aggregated_result)
 
 st.markdown("## Description of models:")