Update app.py

app.py

@@ -10,6 +10,7 @@ from numpy import savetxt, loadtxt
 import numpy as np
 import json
 from copy import deepcopy
+# from sty import fg, bg, ef, rs, RgbBg, Style
 import re
 from tqdm import tqdm
 import gradio as gr
@@ -17,8 +18,8 @@ import gradio as gr
 os.system("python -m spacy download en_core_web_sm")
 nlp = spacy.load("en_core_web_sm")
 tokenizer = RobertaTokenizer.from_pretrained("roberta-base")  
-clause_model = AutoModelForTokenClassification.from_pretrained("./clause_model_512", num_labels=3)
-classification_model = RobertaForSequenceClassification.from_pretrained("./classfication_model", num_labels=18)
+clause_model = AutoModelForTokenClassification.from_pretrained("C:\\Users\\pixin\\Desktop\\Reddit ML\\Trained Models\\clause_model_512", num_labels=3)
+classification_model = RobertaForSequenceClassification.from_pretrained("C:\\Users\\pixin\\Desktop\\Reddit ML\\Trained Models\\classfication_model", num_labels=18)
 
 
 labels2attrs = {
@@ -103,41 +104,43 @@ def seg_clause(text):
         segmented_clauses.append(deepcopy(current_clause))
     return [" ".join(clause) for clause in segmented_clauses if clause is not None]
 
-def pretty_print_segmented_clause(segmented_clauses):
-    np.random.seed(42)
-    bg.orange = Style(RgbBg(255, 150, 50))
-    bg.purple = Style(RgbBg(180, 130, 225))
-    colors = [bg.red, bg.orange, bg.yellow, bg.green, bg.blue, bg.purple]
-    prev_color = 0
-    to_print = []
-    for cl in segmented_clauses:
-        color_choice = np.random.choice(np.delete(np.arange(len(colors)), prev_color))
-        prev_color = color_choice
-        colored_cl = colors[color_choice] + cl + bg.rs
-        to_print.append(colored_cl)
-    print(*to_print, sep=" ")
+# def pretty_print_segmented_clause(segmented_clauses):
+#     np.random.seed(42)
+#     bg.orange = Style(RgbBg(255, 150, 50))
+#     bg.purple = Style(RgbBg(180, 130, 225))
+#     colors = [bg.red, bg.orange, bg.yellow, bg.green, bg.blue, bg.purple]
+#     prev_color = 0
+#     to_print = []
+#     for cl in segmented_clauses:
+#         color_choice = np.random.choice(np.delete(np.arange(len(colors)), prev_color))
+#         prev_color = color_choice
+#         colored_cl = colors[color_choice] + cl + bg.rs
+#         to_print.append(colored_cl)
+#     print(*to_print, sep=" ")
     
 
 def get_pred_classification_labels(clauses, batch_size=32):
     clause2labels = []
-    for i in range(0, len(clauses) + 1, batch_size):
+    for i in range(0, len(clauses), batch_size):
         batch_examples = clauses[i : i + batch_size]
         model_inputs = tokenizer(batch_examples, padding='max_length', max_length=128, truncation=True, return_tensors='pt')
         logits = classification_model(**model_inputs)[0]
         pred_labels = logits.argmax(-1).numpy()
         pred_labels = [index2label[l] for l in pred_labels]
-        
-        clause2labels.extend([(s, str(l),) for s,l in zip(batch_examples, pred_labels)])
+        clause2labels.extend([(s, str(labels2attrs[l]),) for s,l in zip(batch_examples, pred_labels)])
+        print(clause2labels)
     return clause2labels
 
 
 
 def run_pipeline(text):
     snippets = auto_split(text)
+    print(snippets)
     all_clauses = []
     for s in snippets:
         segmented_clauses = seg_clause(s)
         all_clauses.extend(segmented_clauses)
+
     clause2labels = get_pred_classification_labels(all_clauses)
     output_clauses = [(c, str(i + 1))  for i, c in enumerate(all_clauses)]
     return output_clauses, clause2labels
@@ -160,7 +163,7 @@ def run_pipeline(text):
 
 color_panel_1 =  ["red", "green",  "yellow", "DodgerBlue", "orange", "DarkSalmon", "pink", "cyan", "gold", "aqua", "violet"]
 index_colormap = {str(i) : color_panel_1[i % len(color_panel_1)] for i in np.arange(1, 100000)}
-color_panel_2 = ["Violet", "DodgerBlue", "Wheat", "OliveDrab", "DarkKhaki", "DarkSalmon", "Orange", "Gold", "Aqua", "Tomato", "Gray"]
+color_panel_2 = ["Violet", "Gray", "Wheat", "OliveDrab", "DarkKhaki", "DarkSalmon", "Orange", "Gold", "Aqua", "Tomato", "DodgerBlue"]
 str_attrs = [str(v) for v in set(labels2attrs.values())]
 print(str_attrs, len(str_attrs), len(color_panel_2))
 assert len(str_attrs) == len(color_panel_2)