Update app.py

app.py

@@ -3,55 +3,33 @@ import numpy as np
 from huggingface_hub import hf_hub_download
 import os
 
-    
-def extract_spannet_scores(model,input_sentence,length, pos_col, task_col):
-
-    all_scores = []
-    
-    # model = SpanNet.load_model(model_path=path)
-    scores = []
-    model.eval()
-        
-    out_dict = model(sentences=[input_sentence], output_span_scores=True)
-    scores.extend([[t.tolist() for t in o[:l]] for o, l in zip(out_dict['span_scores'], length)])
-    all_scores.append(scores)
-    return all_scores
-
-def pool_span_scores(score_dicts, sent_lens):
-    TAGS = ['B', 'I', 'O']
-    pooled_scores = [[np.argmax([sum([sd[sent_id][token_id][score_id] for sd in score_dicts])
-                                    for score_id in range(len(score_dicts[0][sent_id][token_id]))])
-                                    for token_id in range(sent_lens[sent_id])]
-                        for sent_id in range(len(sent_lens))]
-                        
-    r = [[TAGS[ps] for ps in sent_ps] for sent_ps in pooled_scores]
-    return r 
 
 def predict_label(text):
-    # model_path = 'models/span.model'
+    
     ip = text.split()
     ip_len = [len(ip)]
-    scores = extract_spannet_scores(model,ip,ip_len, pos_col=1, task_col=2)
+    scores = extract_spannet_scores(span_model,ip,ip_len, pos_col=1, task_col=2)
     pooled_scores = pool_span_scores(scores, ip_len)
-    # output=''
-    # for op in pooled_scores[0]:
-    #     output+= op + ','
+    
     return pooled_scores
 
 
 if __name__ == '__main__':
+    
     space_key = os.environ.get('key')
-    filenames = ['network.py', 'layers.py', 'utils.py', 'representation.py']
+    filenames = ['network.py', 'layers.py', 'utils.py', 'representation.py', 'predict.py']
     for file in filenames:
         hf_hub_download('nehalelkaref/stagedNER',
                     filename=file,
                     local_dir='src',
                     token=space_key)
     
-    from src.network import SpanNet
+    from src.predict import extract_spannet_scores,pool_span_scores
+    from src.network import SpanNet, EntNet
     
-    model_path = 'models/span.model'
-    model = SpanNet.load_model(model_path)
+    span_path = 'models/span.model'
+    # span_msa_path = 'models/sp'
+    span_model = SpanNet.load_model(model_path)
     
     iface = gr.Interface(fn=predict_label, inputs="text", outputs="text")
     iface.launch(show_api=False)