mbastardi24/token_classification

1from transformers import pipeline2import gradio as gr3 4token_classification = pipeline("token-classification", model="mbastardi24/roBERTa-finetuned-wnut2017")5 6def post_process(output):7    start_word = 08    end_word =09    current_label = None10    current_word = None11    words = []12 13    for item in output:14      prefix, label = item['entity'][:1],item['entity'][2:]15      if prefix == 'B':16 17        if start_word != end_word != 0:18          words.append({'word':current_word, 'start': start_word, 'end': end_word, 'entity': current_label})19 20        start_word = item['start']21        end_word = item['end']22        current_label = label23        current_word = item['word'][1:]24 25      if prefix == 'I':26        end_word = item['end']27        if item['word'][0] == 'Ġ':28          current_word+=item['word'].replace('Ġ', " ", 1)29        else:30          current_word+=item['word']31 32    words.append({'word':current_word, 'start': start_word, 'end': end_word, 'entity': current_label})33    return words34 35def grad_func(text):36  output = token_classification(text)37  input = post_process(output)38  return {"text": text, "entities": input}39 40token_classification_gradio = gr.Interface(grad_func,41             gr.Textbox(placeholder="Enter sentence here..."),42             gr.HighlightedText())43 44token_classification_gradio.launch()