Spaces:
Runtime error
Runtime error
from transformers import pipeline | |
import gradio as gr | |
pipelines_text = { | |
'Spam': {'BERT': pipeline("text-classification", model="mariagrandury/distilbert-base-uncased-finetuned-sms-spam-detection"), | |
'RoBERTa': pipeline("text-classification", model="mariagrandury/roberta-base-finetuned-sms-spam-detection") | |
}, | |
'Sentiment': { | |
'BERT': pipeline("text-classification", model="lxyuan/distilbert-base-multilingual-cased-sentiments-student"), | |
'RoBERTa': pipeline("text-classification", model="cardiffnlp/twitter-roberta-base-sentiment-latest") | |
}, | |
'Emotion': {'BERT': pipeline("text-classification", model="bhadresh-savani/bert-base-go-emotion"), | |
'RoBERTa': pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base") | |
} | |
} | |
def parseImage(file, radio): | |
return file.name | |
max_textboxes = 100 | |
def change_textboxes(n): | |
return [gr.Textbox.update(visible=True, interactive=True)]*n + [gr.Textbox.update(visible=False, interactive=True)]*(max_textboxes-int(n)) | |
def parseText(text_upload_file, delimeter_dropdown): | |
delimeter_mapping = {'New Line': '\n','Tab': '\t','Comma': ','} | |
delimeter = delimeter_mapping[delimeter_dropdown] | |
text_boxes = ['' for i in range(max_textboxes)] | |
with open(text_upload_file.name, 'r') as f: | |
text_upload = f.read() | |
for idx, text in enumerate(text_upload.split(delimeter)): | |
text_boxes[idx] = text | |
return text_boxes | |
def annotateText(text_boxes_slider, annotation_radio, model_dropdown, *text_boxes_texbox): | |
text_boxes_texbox = [text for text in text_boxes_texbox] | |
res_label = ['' for i in range(max_textboxes)] | |
res_score = ['' for i in range(max_textboxes)] | |
# predictions | |
pipe = pipelines_text[annotation_radio][model_dropdown] | |
predictions = pipe([text_boxes_texbox[i] for i in range(text_boxes_slider)]) | |
for idx, pred in enumerate(predictions): | |
# special case for spam (might change later) | |
if annotation_radio == 'Spam': | |
res_label[idx] = 'Not Spam' if pred['label'] == 'LABEL_0' else 'Spam' | |
else: | |
res_label[idx] = pred['label'] | |
res_score[idx] = '{:.2f}'.format(pred['score']) | |
with open('annotations.csv', 'w') as f: | |
f.write('text,annotation,confidence\n') | |
for idx in range(max_textboxes): | |
if text_boxes_texbox[idx]: | |
f.write('{},{},{}\n'.format(text_boxes_texbox[idx], res_label[idx], res_score[idx])) | |
return ['./annotations.csv'] + text_boxes_texbox + res_label + res_score | |
with gr.Blocks() as demo: | |
gr.Markdown("# Data Annotation Tool") | |
gr.Markdown('Upload a file or enter text in the Data Viewer section. Sample files are at the end of the page.') | |
with gr.Tab("Text"): | |
with gr.Row(): | |
with gr.Column(): | |
gr.Markdown("## Data Upload") | |
text_upload_file = gr.File(file_types=['text']) | |
delimeter_dropdown = gr.Dropdown(choices=['New Line','Tab','Comma'], label='Delimeter') | |
text_upload_button = gr.Button('Parse File') | |
with gr.Row(): | |
with gr.Column(): | |
gr.Markdown("## Data Viewer") | |
# slider component | |
text_boxes_slider = gr.Slider(1, max_textboxes, value=3, step=1) | |
# text box components (3 visible and max_textboxes-3 not visible) | |
text_boxes_texbox = [gr.Textbox(show_label=False,interactive=True) for i in range(3)] + [gr.Textbox(show_label=False, visible=False) for i in range(max_textboxes-3)] | |
annotation_radio = gr.Radio(choices=['Spam', 'Sentiment', 'Emotion'], label='Annotation', value='RoBERTa') | |
model_dropdown = gr.Dropdown(choices=['BERT', 'RoBERTa'], label='Model') | |
text_submit_button = gr.Button('Annotate Data') | |
with gr.Row(): | |
gr.Markdown("## Data Output") | |
with gr.Row(): | |
with gr.Column(scale=6): | |
gr.Markdown("Text") | |
text_output_boxes = [gr.Textbox(show_label=False,interactive=False) for i in range(3)] + [gr.Textbox(show_label=False, visible=False, interactive=False) for i in range(max_textboxes-3)] | |
with gr.Column(scale=1): | |
gr.Markdown("Annotations") | |
text_output_annotations_boxes = [gr.Textbox(show_label=False,interactive=False) for i in range(3)] + [gr.Textbox(show_label=False, visible=False, interactive=False) for i in range(max_textboxes-3)] | |
with gr.Column(scale=1): | |
gr.Markdown("Confidence") | |
text_output_confidence_boxes = [gr.Textbox(show_label=False,interactive=False) for i in range(3)] + [gr.Textbox(show_label=False, visible=False, interactive=False) for i in range(max_textboxes-3)] | |
text_ouput_file = gr.File(label='File Output', file_types=['csv']) | |
gr.Markdown("## Test Examples") | |
with gr.Row(): | |
with gr.Column(): | |
gr.Examples( | |
examples=[['./examples/text/spam.txt', 'New Line'],['./examples/text/sentiment.txt', 'New Line'],['./examples/text/emotion.txt', 'New Line']], | |
fn=parseText, | |
inputs=[text_upload_file, delimeter_dropdown], | |
outputs=text_boxes_texbox, | |
cache_examples=True | |
) | |
# event listeners | |
text_upload_button.click(fn=parseText, inputs=[text_upload_file, delimeter_dropdown], outputs=text_boxes_texbox) | |
text_boxes_slider.change(fn=change_textboxes, inputs=text_boxes_slider, outputs=text_boxes_texbox) | |
text_boxes_slider.change(fn=change_textboxes, inputs=text_boxes_slider, outputs=text_output_boxes) | |
text_boxes_slider.change(fn=change_textboxes, inputs=text_boxes_slider, outputs=text_output_annotations_boxes) | |
text_boxes_slider.change(fn=change_textboxes, inputs=text_boxes_slider, outputs=text_output_confidence_boxes) | |
text_submit_button.click(fn=annotateText, inputs=[text_boxes_slider, annotation_radio, model_dropdown] + text_boxes_texbox, outputs=[text_ouput_file]+text_output_boxes + text_output_annotations_boxes+text_output_confidence_boxes) | |
with gr.Tab("Image"): | |
with gr.Row(): | |
gr.Markdown("## Coming Soon!") | |
# with gr.Row(): | |
# file_image = gr.File(file_count=['directory'],file_types=['image'], label='File Upload') | |
# image = gr.Image() | |
# with gr.Row(): | |
# radio_image = gr.Radio(choices=['Object Detection'], label='Annotation') | |
# models_image = gr.Dropdown(choices=['DETR'], label='Model') | |
# with gr.Row(): | |
# button_image = gr.Button('Submit') | |
# with gr.Row(): | |
# output_image = gr.File(label='File Output', file_types=['image']) | |
# image tab event listeners | |
# button_image.click(fn=doImage, inputs=[file_image, radio_image], outputs=output_image) | |
if __name__ == "__main__": | |
demo.launch() |