| import gradio as gr | |
| from transformers import pipeline | |
| raw_model_name = 'distilbert-base-uncased' | |
| raw_model = pipeline('sentiment-analysis', model=raw_model_name) | |
| fine_tuned_model_name = 'distilbert-base-uncased-finetuned-sst-2-english' | |
| fine_tuned_model = pipeline('sentiment-analysis', model=fine_tuned_model_name) | |
| def get_model_output(input_text, model_choice): | |
| raw_result = raw_model(input_text) | |
| fine_tuned_result = fine_tuned_model(input_text) | |
| return format_model_output(raw_result[0]), format_model_output(fine_tuned_result[0]) | |
| def format_model_output(output): | |
| return f"I am {output['score']*100:.2f}% sure that the sentiment is {output['label']}" | |
| iface = gr.Interface( | |
| fn=get_model_output, | |
| title="DistilBERT Sentiment Analysis", | |
| inputs=[ | |
| gr.Textbox(label="Input Text"), | |
| ], | |
| outputs=[ | |
| gr.Textbox(label="Base DistilBERT output (distilbert-base-uncased)"), | |
| gr.Textbox(label="Fine-tuned DistilBERT output") | |
| ], | |
| ) | |
| iface.launch() | |