update app.py to use puberta-news

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 
 MODEL_URL = "https://huggingface.co/dsfsi/PuoBERTa-News"
@@ -8,33 +7,45 @@ 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):
-    # create pipeline
     clasifer = pipeline("sentiment-analysis", tokenizer=tokenizer, model=model, return_all_scores=True)
 
     preds = clasifer(news)
 
-    preds_dict={}
+    preds_dict = {}
     for pred in preds[0]:
-        preds_dict[pred['label']] = pred['score']
+        label = categories.get(pred['label'], pred['label'])
+        preds_dict[label] = pred['score']
 
     return preds_dict
 
-
 gradio_ui = gr.Interface(
     fn=prediction,
     title="Setswana News Classification",
-    description=f"Enter Turkish news article to see the category of the news.\n For this classification, the {MODEL_URL} model was used.",
+    description=f"Enter Setswana news article to see the category of the news.\n For this classification, the {MODEL_URL} model was used.",
     examples=[
-        ['Son dakika haberi... Sağlık Bakanlığı corona virüsü (koronavirüs) için son gelişmeleri ve Türkiye Günlük Korona Tablosu\'nu açıkladı. Türkiye\'de son 24 saatte 355 bin 947 Kovid-19 testi yapıldı, 26 bin 896 kişinin testi pozitif çıktı, 210 kişi yaşamını yitirdi. Tabloda vaka ve vefat sayılarına dikkat çeken Bakan Koca, "Bu ölümlerin bir kısmı birden çok sebebin bir araya gelmesine bağlı ve önlenebilir. Örneğin, Covid-19’a kronik hastalıkların eşlik etmesi riski artıran bir faktör. Tam doz aşı bu riski azaltıyor." çağrısında bulundu. Öte yandan ikinci doz aşılama oranı %75\'e çıkan Tunceli, mavi kategoriye yükseldi.'],
-        ["Galatasaray'dan sezon başında ayrılan ve Rayo Vallecano'nun yolunu tutan Radamel Falcao, İspanya'da şov yapmaya devam ediyor. Kolombiyalı yıldızın, Barcelona karşısında attığı gol maça damga vurdu."],
-        ["Instagram, 3 yıl önce kullanıcıların beğenisine sunduğu IGTV özelliğini kaldırıyor. IGTV’nin yerine gelecek yeni özelliğin ise platformun Reels'tan elde ettiği başarıyı pekiştirip içerik üreticileri için kullanışlı bir zemin oluşturması bekleniyor."],
-        ["Kamu bankaları Halkbank, Vakıfbank ve Ziraat Bankası ortak bir açıklama yayınlayarak kredi faiz oranlarında indirim yapıldığını açıkladı. Söz konusu indirimler ticari ve konut kredileri kapsıyor. Tüketici kredisi kapsam dışında bırakıldı."]
+        ['Ka Letsatsi la Aforika, Aforika Borwa e tla be e keteka mabaka a boikemelo, le diketso tse di siameng tse e di dirileng go tokafatsa dikamano tsa yona le dinaga tse dingwe tsa Aforika.'],
+        ["Thuto ya Setswana ke nngwe ya dithuto tse di botlhokwa mo sekolong se se tlhamaletseng go ruta bana ba ba mo lefatsheng la Botswana."],
+        ["Mo kgweding e e fetileng, dipuisano tsa ditheko tsa dijalo di ile tsa tswelela, ka batho ba rekang le barui ba ba ruileng."],
+        ["Masole a Aforika Borwa a ne a ya kwa Mozambique go tlisetsa motlakase morago ga maduo a kgatlha."],
     ],
-    inputs=gr.inputs.Textbox(lines=10, label="Paste some news here"),
+    inputs=gr.inputs.Textbox(lines=10, label="Paste some Setswana news here"),
     outputs=gr.outputs.Label(num_top_classes=5, type="auto", label="News categories probabilities"),
     theme="huggingface",
-    article="<p style='text-align: center'>for our other AI works: <a href='https://www.kodiks.com/ai_solutions.html' target='_blank'>https://www.kodiks.com/ai_solutions.html</a> | <a href='https://twitter.com/KodiksBilisim' target='_blank'>Contact us</a></p>",
+    article="<p style='text-align: center'>For our other AI works: <a href='https://www.kodiks.com/ai_solutions.html' target='_blank'>https://www.kodiks.com/ai_solutions.html</a> | <a href='https://twitter.com/KodiksBilisim' target='_blank'>Contact us</a></p>",
 )
 
-gradio_ui.launch()
+gradio_ui.launch()