updated

app.py

@@ -8,7 +8,7 @@ sys.path.append(project_dir)
 
 from model import MultiTaskBertModel
 from data_loader import load_dataset
-from utils import bert_config, tokenizer, intent_ids_to_labels, intent_labels_to_ids
+from utils import bert_config, tokenizer, intent_ids_to_labels, intent_labels_to_ids, ner_labels_to_ids, ner_ids_to_labels
 
 config = bert_config()
 dataset = load_dataset("training_dataset")
@@ -18,6 +18,9 @@ model.load_state_dict(torch.load("pytorch_model.bin"))
 
 model.eval()
 
+ner_label_to_id = ner_labels_to_ids()
+ner_id_to_label = ner_ids_to_labels(ner_label_to_id)
+
 def predict(input_data):
 
     tok = tokenizer()
@@ -49,16 +52,22 @@ def predict(input_data):
         if not word.strip():
             continue
 
-        aligned_predictions.append((word, int(prediction)))
+        aligned_predictions.append((word, ner_id_to_label[int(prediction)]))
 
     labels = intent_labels_to_ids()
     intent_labels = intent_ids_to_labels(labels)
+    intent_labels = intent_labels[int(intent_logits)]
 
     return f"Ner logits: {aligned_predictions}, Intent Label: {intent_labels}"
 
 title = "Multi Task Model"
 description = '''
-The model was trained to do NER and Intent classification for a scheduler
+This model is designed for a scheduler application, capable of handling various tasks such as setting 
+timers, scheduling meetings, appointments, and alarms. It provides Named Entity Recognition (NER) labels 
+to identify specific entities within the input text, along with an Intent label to determine the 
+overall task intention. The model's outputs facilitate efficient task management and organization, 
+enabling seamless interaction with the scheduler application.
+<img src="bart.jpg" width=300px>
 '''
 
 gr.Interface(