Spaces:
Running
Running
| import gradio as gr | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline | |
| import pandas as pd | |
| MODEL_URL = "https://huggingface.co/dsfsi/PuoBERTa-News" | |
| WEBSITE_URL = "https://www.kodiks.com/ai_solutions.html" | |
| tokenizer = AutoTokenizer.from_pretrained("dsfsi/PuoBERTa-News") | |
| model = AutoModelForSequenceClassification.from_pretrained("dsfsi/PuoBERTa-News") | |
| categories = { | |
| "arts_culture_entertainment_and_media": "Botsweretshi, setso, boitapoloso le bobegakgang", | |
| "crime_law_and_justice": "Bosenyi, molao le bosiamisi", | |
| "disaster_accident_and_emergency_incident": "Masetlapelo, kotsi le tiragalo ya maemo a tshoganyetso", | |
| "economy_business_and_finance": "Ikonomi, tsa kgwebo le tsa ditšhelete", | |
| "education": "Thuto", | |
| "environment": "Tikologo", | |
| "health": "Boitekanelo", | |
| "politics": "Dipolotiki", | |
| "religion_and_belief": "Bodumedi le tumelo", | |
| "society": "Setšhaba" | |
| } | |
| def prediction(news): | |
| clasifer = pipeline("sentiment-analysis", tokenizer=tokenizer, model=model, return_all_scores=True) | |
| preds = clasifer(news) | |
| preds_dict = {categories.get(pred['label'], pred['label']): pred['score'] for pred in preds[0]} | |
| return preds_dict | |
| def file_prediction(file): | |
| # Load the file (CSV or text) | |
| if file.name.endswith('.csv'): | |
| df = pd.read_csv(file.name) | |
| news_list = df.iloc[:, 0].tolist() | |
| else: | |
| news_list = [file.read().decode('utf-8')] | |
| results = [] | |
| for news in news_list: | |
| results.append(prediction(news)) | |
| return results | |
| gradio_ui = gr.Interface( | |
| fn=prediction, | |
| title="Setswana News Classification", | |
| description=f"Enter Setswana news article to see the category of the news.\n For this classification, the {MODEL_URL} model was used.", | |
| inputs=gr.Textbox(lines=10, label="Paste some Setswana news here"), | |
| outputs=gr.Label(num_top_classes=5, label="News categories probabilities"), | |
| ) | |
| gradio_file_ui = gr.Interface( | |
| fn=file_prediction, | |
| title="Upload File for Setswana News Classification", | |
| description=f"Upload a text or CSV file with Setswana news articles. The first column in the CSV should contain the news text.", | |
| inputs=gr.File(label="Upload text or CSV file"), | |
| outputs=gr.Dataframe(headers=["News Text", "Category Predictions"], label="Predictions from file"), | |
| ) | |
| gradio_combined_ui = gr.TabbedInterface([gradio_ui, gradio_file_ui], ["Text Input", "File Upload"]) | |
| css = """ | |
| body { | |
| background-color: white !important; | |
| color: black !important; | |
| } | |
| .gradio-container { | |
| background-color: white !important; | |
| color: black !important; | |
| } | |
| .gr-input, .gr-button, .gr-textbox, .gr-file, .gr-dataframe { | |
| background-color: white !important; | |
| color: black !important; | |
| border-color: #ccc !important; | |
| } | |
| .gr-button { | |
| background-color: #f0f0f0 !important; | |
| color: black !important; | |
| border: 1px solid #ccc !important; | |
| } | |
| .gr-dataframe th, .gr-dataframe td { | |
| background-color: #f9f9f9 !important; | |
| color: black !important; | |
| } | |
| """ | |
| gradio_combined_ui.launch(css=css) | |