Enhanced visuals

app.py

@@ -28,10 +28,10 @@ model_mapping = {
 
 # Example text included within the interface
 exampleText = [
-    ["ex1"],
-    ["ex2"],
-    ["ex3"],
-    ["ex4"]
+    "ex1",
+    "ex2",
+    "ex3",
+    "ex4"
 ]
 
 # Example models and datasets included within the interface
@@ -47,7 +47,7 @@ exampleDatasets = ["No Dataset Finetuning",
 examples = [[random.choice(exampleModels), random.choice(exampleDatasets), random.choice(exampleText)] for example in exampleText]
 
 
-def detect_ai_generated_text(model: str, dataset: str, text: str) -> str:
+def detect_ai_generated_text(model: str, dataset: str, text: str) -> dict:
     # Get the fine-tuned model using mapping
     finetuned_model = model_mapping.get((model, dataset))
 
@@ -63,7 +63,13 @@ def detect_ai_generated_text(model: str, dataset: str, text: str) -> str:
     label = "AI-generated" if result[0]['label'] == 'LABEL_1' else "Human-written"
     score = result[0]['score']
 
-    return f"{label} with confidence {score * 100:.2f}%"
+    # Create HTML for the colored bars
+    ai_score = score if label == "AI-generated" else 1 - score
+    human_score = 1 - ai_score
+    ai_bar = f'<div style="background-color: red; width: {ai_score * 100}%; height: 20px;"></div>'
+    human_bar = f'<div style="background-color: blue; width: {human_score * 100}%; height: 20px;"></div>'
+
+    return {"label": f"{label} with confidence {score * 100:.2f}%", "ai_bar": ai_bar, "human_bar": human_bar}
 
 
 interface = gr.Interface(
@@ -73,7 +79,11 @@ interface = gr.Interface(
         gr.Dropdown(choices=datasets, label="Dataset"),
         gr.Textbox(lines=5, label="Input Text")
     ],
-    outputs=gr.Textbox(label="Output"),
+    outputs=[
+        gr.Label(label="Output"),
+        gr.HTML(label="AI-generated"),
+        gr.HTML(label="Human-written")
+    ],
     examples=examples,
     title="AI Generated Text Detection"
 )