update interface to be consistent and allow file uploads

app.py

@@ -1,5 +1,6 @@
 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"
@@ -22,31 +23,55 @@ categories = {
 
 def prediction(news):
     clasifer = pipeline("sentiment-analysis", tokenizer=tokenizer, model=model, return_all_scores=True)
-
     preds = clasifer(news)
-
-    preds_dict = {}
-    for pred in preds[0]:
-        label = categories.get(pred['label'], pred['label'])
-        preds_dict[label] = pred['score']
-
+    preds_dict = {categories.get(pred['label'], pred['label']): pred['score'] for pred in preds[0]}
     return preds_dict
 
+def file_prediction(file):
+    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')]  # Load plain text
+    
+    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.",
-    examples=[
-        ['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.Textbox(lines=10, label="Paste some Setswana news here"),
     outputs=gr.Label(num_top_classes=5, 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>",
+    theme="default",
+    css="""
+    body {
+        background-color: white !important;
+        color: black !important;
+    }
+    .gradio-container {
+        background-color: white !important;
+        color: black !important;
+    }
+    .gr-button {
+        background-color: #f0f0f0 !important;
+        color: black !important;
+    }
+    """
+)
+
+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"),
+    theme="default"
 )
 
-gradio_ui.launch()
+gradio_combined_ui = gr.TabbedInterface([gradio_ui, gradio_file_ui], ["Text Input", "File Upload"])
 
+gradio_combined_ui.launch()