import gradio as gr
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)
    pooled_scores = pool_span_scores(scores, ip_len)
    output=''
    for op in pooled_scores[0]:
        output+= op + ','
    print('OUTPUT HERE')
    return 'output'
def temp(text):
    print('IN FUNCTION')
    return text
print('STARTING ..')


if __name__ == '__main__':
    space_key = os.environ.get('key')
    filenames = ['network.py', 'layers.py', 'utils.py', 'representation.py']
    for file in filenames:
        hf_hub_download('nehalelkaref/stagedNER',
                    filename=file,
                    local_dir='src',
                    token=space_key)
    
    from src.network import SpanNet
    
    model_path = 'models/span.model'
    model = SpanNet.load_model(model_path)
    
    iface = gr.Interface(fn=predict_label, inputs="text", outputs="text")
    iface.launch(show_api=False)