import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download
import os


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


if __name__ == '__main__':
    
    space_key = os.environ.get('key')
    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.predict import extract_spannet_scores,pool_span_scores
    from src.network import SpanNet, EntNet
    
    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)