Update app.py

app.py

@@ -8,10 +8,14 @@ def predict_label(text):
     
     ip = text.split()
     ip_len = [len(ip)]
-    scores = extract_spannet_scores(span_model,ip,ip_len, pos_col=1, task_col=2)
-    pooled_scores = pool_span_scores(scores, ip_len)
     
-    return pooled_scores
+    span_scores = extract_spannet_scores(span_model,ip,ip_len, pos_col=1, task_col=2)
+    pooled_scores = pool_span_scores(span_scores, ip_len)
+    
+    ent_scores = extract_ent_scores(ent_model, ip,pooled_scores, pos_col=1, task_col=2)
+    combined_sequences, ent_pred_tags = pool_ent_scores(ent_scores, ip_len)
+    
+    return combined_sequences
 
 
 if __name__ == '__main__':
@@ -30,8 +34,10 @@ if __name__ == '__main__':
     from src.network import SpanNet, EntNet
     
     span_path = 'models/span.model'
+    entity_path= 'models/entity.msa.model'
     # span_msa_path = 'models/sp'
     span_model = SpanNet.load_model(span_path)
+    entity_model = EntNet.load_model(entity_path)
     
     iface = gr.Interface(fn=predict_label, inputs="text", outputs="text")
     iface.launch(show_api=False)